Context
The current continuous-batching rollout path assumes equal prompt widths and
uses a shared logical decode position for active rows. This keeps the initial
prototype simple, but it can require a larger cache buffer than necessary when
prompt lengths vary.
For rollout workloads, we do not need to preserve serving-oriented TTFT/TPOT
ordering. This gives us an opportunity to reorder prompts internally and
reduce unused cache space.
Proposed direction
- Sort prompts from longest to shortest internally and right-align the initial
prompts so that their decode fronts are aligned.
- When a row retires, admit a pending prompt only when its prefill can fit
behind the current decode front.
- Track each active row's logical span start and position independently.
- Periodically detect dead zones before the earliest active span and trim them
by shifting active KV rows, attention metadata, and write positions to the
left.
- Restore the original request-row order in the returned
RolloutBatch.
Correctness requirements
- Support mixed prompt widths.
- Preserve token-level equality with the sequential eager reference.
- Cover EOS retirement and prompt refill.
- Verify KV-cache contents, attention masks, logical positions, and write
positions after trimming.
- Preserve the original request-row order in the output.
Performance requirements
Report:
- peak cache length;
- cache memory usage;
- bytes moved during trimming;
- trimming frequency and latency;
- end-to-end rollout throughput.
The cache capacity should be bounded by the maximum prompt length plus the
shared generation budget, subject to the selected scheduling policy.
Scope
Follow-up to DeepSpeed PR #8368. The current equal-width StaticCache path remains
the baseline until this optimization is implemented and validated.
cc @delock Appreciate your insights. Happy to hear any feedback or recommendations.
Context
The current continuous-batching rollout path assumes equal prompt widths and
uses a shared logical decode position for active rows. This keeps the initial
prototype simple, but it can require a larger cache buffer than necessary when
prompt lengths vary.
For rollout workloads, we do not need to preserve serving-oriented TTFT/TPOT
ordering. This gives us an opportunity to reorder prompts internally and
reduce unused cache space.
Proposed direction
prompts so that their decode fronts are aligned.
behind the current decode front.
by shifting active KV rows, attention metadata, and write positions to the
left.
RolloutBatch.Correctness requirements
positions after trimming.
Performance requirements
Report:
The cache capacity should be bounded by the maximum prompt length plus the
shared generation budget, subject to the selected scheduling policy.
Scope
Follow-up to DeepSpeed PR #8368. The current equal-width StaticCache path remains
the baseline until this optimization is implemented and validated.
cc @delock Appreciate your insights. Happy to hear any feedback or recommendations.