[liuhao-labs]

Summary

Max-length filtering already tokenizes every retained text prompt, but default rollout preparation encodes the same rendered prompt again. Reuse those IDs for stock training and evaluation, eliminating the second encoding pass.

The cache contains only immutable packed uint32 IDs (4 bytes per token). Deep-copied siblings share the payload and each materializes its own mutable token list on consumption. The cache is excluded from serialization and dataclass reconstruction. Multimodal outputs and custom rollout/generate/data-source contracts remain uncached.

Training is wired through the current slime.data.data_source API. Missing rollout configuration disables caching; existing tokens and multimodal processing retain their precedence.

Related to #1111.

Validation

• 111 local CPU tests passed across filtering, Sample serialization/reconstruction, train/eval cache ownership, buffer/checkpoint behavior, fully async rollout, and plugin contracts.
• Direct CI-style filtering runner: 25 passed.
• Repository-pinned pre-commit checks passed.
• On the original 10,000-prompt in-memory BPE benchmark, second-pass tokenizer.encode calls dropped from 10,000 to 0. The immediate same-process comparison for filtering plus preparation was 2.65s versus 0.57s; packed payload size was 8.35 MiB. Timing is workload-dependent.
• The updated GitHub CPU/agent matrix also exercises distributed Ray/Straw fan-out against the current upstream data-source API.

[nanjiangwill]

i checked this before and decided not to add it because re-tokenization especially in text-only case wont add too much latency. theoretically yes we should have this.

[liuhao-labs]

Thanks for the context. I measured the retained text-only path before opening this, because the benefit depends on dataset size and prompt length.

On the same local workload (10,000 retained prompts, 2,189,001 prompt token IDs, an in-memory ByteLevel BPE), the current path made one batched tokenization call during filtering plus 10,000 per-sample encode() calls during rollout preparation. With this change, the second count is zero.

The immediate same-process comparison was:

• current path: 2.652 s total (0.586 s filter + 2.066 s prepare)
• cached path: 0.571 s total (0.541 s filter + 0.030 s prepare)
• rollout preparation: 98.6% lower, total for the two stages: 78.5% lower

The retained cache payload was 8,756,004 bytes (8.35 MiB) as packed uint32 IDs. It is only created when max-length filtering has already tokenized the prompt, only for retained text prompts on the stock rollout path, and it is consumed before Ray transfer. Multimodal processor outputs and custom rollout/data-source paths remain uncached.

So this is not intended to add tokenization work when filtering is disabled; it only reuses work that has already happened. The exact speedup will vary with tokenizer and prompts, but the invariant improvement is one fewer encoding pass per retained text prompt.

[liuhao-labs]

Rebased onto current main (3778dbf) and preserved the existing implementation unchanged. The stale baseline failure from the removed test_rollout_validation.py is no longer part of the new CI run. Local checks passed: test_filter_long_prompt.py (22), test_sample.py (20), test_rollout_data_utils.py (6), plus the repository-pinned format/lint checks. @nanjiangwill, this is ready for another look when convenient.

[liuhao-labs]

@nanjiangwill, refreshed this against the current v0.4 mainline. The cache is wired through the relocated slime.data.data_source API, and missing rollout configuration disables it. The transient payload now lives outside the dataclass fields, preserving the new fan-out path that reconstructs Samples from dataclasses.fields() while retaining immutable sharing across deep-copied siblings.

Local validation: 111 CPU tests passed, including stock data sources, repeated epochs/siblings, serialization/reconstruction, eval, buffer/checkpoint behavior, fully async rollout, and plugin contracts. The full refreshed PR Test run is green with 89 successful jobs and no failures. Filtering, Sample, buffer, and real distributed Ray/Straw fan-out all passed, as did pre-commit.

The PR is mergeable and ready for another look. The workload benchmark and bounded text-only cache scope are documented in the description.