Skip to content

[FEATURE] Align decode fronts and trim dead zones in continuous batching #8462

Description

@nathon-lee

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

  1. Sort prompts from longest to shortest internally and right-align the initial
    prompts so that their decode fronts are aligned.
  2. When a row retires, admit a pending prompt only when its prefill can fit
    behind the current decode front.
  3. Track each active row's logical span start and position independently.
  4. 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.
  5. 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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions