deepspeedai / deepspeedai/DeepSpeed
[FEATURE] Align decode fronts and trim dead zones in continuous batching
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the current equal-width StaticCache path and the rollout changes from DeepSpeed PR #8368. Trace how RolloutBatch ordering, KV-cache contents, attention masks, logical positions, and write positions are handled during prompt refill and EOS retirement. Done means mixed prompt widths preserve sequential eager token equality, trim dead zones correctly, restore request order, and report the requested cache and throughput metrics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100