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

[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.