[Scheduler] Shortest-job-first prefill admission to cut p90 TTFT

[ganyi1996ppo]

Under FCFS a short request queued behind a long prompt pays that prompt's entire prefill before it sees a first token. --scheduling- policy sjf orders the waiting queue by remaining prefill work instead, so the many short requests clear first and only the few long ones pay. Under disaggregated prefill it also hands short requests to the decode node's batch sooner.

The sort key has two tiers. The front tier holds requests that must not be reordered by length: a preemption victim, which already spent a forward pass and had its KV thrown away, and any request skipped --sjf-max-skip-steps times, which bounds how long the policy may starve a long prompt (0 removes the bound entirely). Everything else ranks by uncached token count. Both tiers sort stably, so the policy degrades to FCFS on a uniform workload.

Two ordering constraints the implementation has to respect:

• The sort runs before _promote_ready_remote_kv_requests and _park_ready_offload_partial_prefills, whose front-of-queue placements are correctness, not preference, and must survive it.
• preempt can no longer rely on appendleft for head position, since the next step re-sorts. It sets is_preempted instead, cleared on admission -- a one-shot retry, not a permanent promotion.

Skips are charged only on steps that actually admitted a prefill. On a step where the engine admitted nobody, no request was passed over in favour of another, and counting it would let a KV stall age the whole queue into the front tier at once.

Known limitation, documented at the call site and in the scheduling guide: num_cached_tokens is written by postprocess and the offload loader, not by the admission-loop prefix probe (which writes prefix_cache_hit_tokens), and deallocate zeroes it on preemption. So for a request that has not yet run a prefill chunk the key equals raw prompt length, and a long-but-cache-hot prompt is ordered as if it were cold. Making the key cache-aware means probing per waiting request per step -- a real cost that deserves its own measurement.

Default stays "fcfs"; this is opt-in. Head-of-line blocking in the admission loop is unchanged and out of scope.

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2158 --add-label <label>

[Copilot]

🟡 Changes recommended

PrefillScheduler’s SJF bookkeeping should clear promotion state on admission, and sjf_max_skip_steps needs explicit non-negative validation to avoid silently disabling the starvation bound.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds an opt-in shortest-job-first (SJF) admission policy for prefill scheduling to reduce head-of-line blocking and improve p90 TTFT by prioritizing requests with less remaining prefill work. The change integrates across colocated and disaggregated (PrefillScheduler) paths, exposes CLI/config knobs, and documents the policy and its known limitations.

Changes:

• Implement scheduling_policy={"fcfs","sjf"} with shared validation + shared SJF sort key, plus skip-based starvation bounding (sjf_max_skip_steps) and preemption-victim protection.
• Wire the policy through config + CLI, and document behavior/limitations in scheduling and configuration guides.
• Add targeted unit tests for ordering, stability, starvation bound behavior, preemption retry behavior, and validation.

File summaries

File	Description
atom/model_engine/scheduler.py	Core SJF policy implementation for colocated and disaggregated schedulers, incl. validation and queue ordering/aging hooks.
atom/model_engine/sequence.py	Adds per-sequence SJF bookkeeping fields (num_skipped_steps, is_preempted).
atom/model_engine/arg_utils.py	Exposes --scheduling-policy and --sjf-max-skip-steps CLI flags and EngineArgs fields.
atom/config.py	Adds Config defaults and documentation comments for the new scheduling knobs.
docs/scheduling_kv_cache_guide.md	Documents SJF behavior, ordering tiers, skip charging, and cache-awareness limitation.
docs/configuration_guide.md	Adds scheduling_policy / sjf_max_skip_steps to the configuration reference.
tests/conftest.py	Extends MockConfig defaults to include the new config fields.
tests/test_scheduler.py	Adds extensive scheduler-level tests covering SJF ordering, stability, skip aging rules, and policy validation.
tests/test_prefill_scheduler.py	Adds PrefillScheduler SJF ordering + starvation bound tests for disaggregated prefill.
tests/test_arg_utils_spec.py	Adds CLI parsing tests ensuring scheduling-policy flags reach engine kwargs and invalid values fail at parse time.

Review details

Suppressed comments (1)

atom/model_engine/scheduler.py:2847

• Same as the colocated Scheduler: PrefillScheduler should reject negative sjf_max_skip_steps values, otherwise a typo like -1 silently removes the starvation bound.

        self.scheduling_policy = validate_scheduling_policy(config.scheduling_policy)
        self.sjf_max_skip_steps = config.sjf_max_skip_steps

• Files reviewed: 10/10 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

It needs small but important correctness/robustness fixes (parameter validation for sjf_max_skip_steps and consistent clearing of SJF promotion flags on PrefillScheduler admission) before it can be safely approved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

atom/model_engine/scheduler.py:557

• sjf_max_skip_steps is accepted from config/CLI without validation; negative values would silently disable promotion and make the starvation bound behave unexpectedly. Since docs/CLI treat 0 as the only “unbounded” value, consider rejecting < 0 early (alongside the policy validation).

        self.scheduling_policy = validate_scheduling_policy(config.scheduling_policy)
        self.sjf_max_skip_steps = config.sjf_max_skip_steps

atom/model_engine/scheduler.py:2962

• SJF skip accounting relies on Sequence.num_skipped_steps/is_preempted being “one-shot” promotion signals, but PrefillScheduler._schedule never clears them when a sequence is admitted (unlike Scheduler._schedule_prefill_seq). If a sequence is ever re-queued (e.g. retries/error paths), it could remain permanently in the front tier.

            # Only the ready set can be "passed over" -- a sequence still
            # awaiting its BlockAssignment was never a candidate this step.
            if self.scheduling_policy == "sjf" and num_seqs:
                for seq in ready:
                    if seq.id not in scheduled_seqs:

• Files reviewed: 11/11 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Operator-facing sjf_max_skip_steps is currently not validated (negative values silently disable the bound), so misconfiguration can produce surprising behavior instead of failing fast.

Review details

Suppressed comments (2)

atom/model_engine/scheduler.py:595

• sjf_max_skip_steps is used as a fairness/ordering guard but is never validated; negative values currently disable the bound implicitly (because 0 < max_skip_steps is false), which is surprising and inconsistent with the documented semantics (0 is the explicit ‘unbounded’ setting). Consider rejecting negative values early so misconfiguration fails loudly.

        self.enable_chunked_prefill = config.enable_chunked_prefill
        self.scheduling_policy = validate_scheduling_policy(config.scheduling_policy)
        self.sjf_max_skip_steps = config.sjf_max_skip_steps

atom/model_engine/scheduler.py:3425

• Same as the colocated scheduler: sjf_max_skip_steps is not validated here, so a negative value silently disables the starvation bound. Since this is an operator-facing knob, it’s safer to fail fast on invalid values.

        self.scheduling_policy = validate_scheduling_policy(config.scheduling_policy)
        self.sjf_max_skip_steps = config.sjf_max_skip_steps

• Files reviewed: 12/12 changed files
• Comments generated: 0 new
• Review effort level: Lite