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

atom/config.py

@@ -1569,6 +1569,17 @@ class Config:
     # `BlockManager._record_checkpoint_demand` for the placement.
     state_checkpoint_demand: bool = True
     scheduler_delay_factor: float = 0.0
+    # Order in which waiting requests are admitted to prefill.
+    #   "fcfs"  arrival order (the historical behaviour)
+    #   "sjf"   shortest-job-first: fewest uncached prompt tokens first, which
+    #           lowers p90 TTFT at the cost of the longest requests' TTFT.
+    # See Scheduler._reorder_waiting_shortest_first.
+    scheduling_policy: str = "fcfs"
+    # Under "sjf", a waiting request skipped this many scheduling steps is
+    # promoted to the front regardless of length, bounding the starvation the
+    # policy would otherwise allow. 0 removes the bound (pure shortest-job-
+    # first). Ignored under "fcfs".
+    sjf_max_skip_steps: int = 64
     max_num_seqs: int = 512
     max_model_len: int | None = None
     gpu_memory_utilization: float = 0.9

atom/model_engine/arg_utils.py

@@ -64,6 +64,8 @@ class EngineArgs:
     throughput_log_interval: float = 10.0
     cache_hit_rate_window: int = 1000
     scheduler_delay_factor: float = 0.0
+    scheduling_policy: str = "fcfs"
+    sjf_max_skip_steps: int = 64
     max_num_seqs: int = 512
     gpu_memory_utilization: float = 0.9
     cudagraph_capture_sizes: str = "[1,2,4,8,16,32,48,64,128,256]"
@@ -399,6 +401,32 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
                 "prefills with decode for lower ITL."
             ),
         )
+        parser.add_argument(
+            "--scheduling-policy",
+            type=str,
+            default="fcfs",
+            choices=["fcfs", "sjf"],
+            help=(
+                "Order in which waiting requests are admitted to prefill. "
+                "'fcfs' is arrival order. 'sjf' is shortest-job-first: the "
+                "request with the fewest remaining prefill tokens goes first, "
+                "which lowers p90 TTFT and raises TTFT for the longest "
+                "prompts. Under disaggregated prefill it also hands short "
+                "requests to the decode node's batch sooner."
+            ),
+        )
+        parser.add_argument(
+            "--sjf-max-skip-steps",
+            type=int,
+            default=64,
+            help=(
+                "Starvation bound for --scheduling-policy sjf: a request "
+                "passed over this many scheduling steps is admitted next "
+                "regardless of length. 0 removes the bound entirely (pure "
+                "shortest-job-first, in which a long prompt can be starved "
+                "indefinitely). Ignored under 'fcfs'."
+            ),
+        )
         parser.add_argument(
             "--attn-prefill-chunk-size",
             type=int,

atom/model_engine/block_manager.py

@@ -846,6 +846,40 @@ def _chain_to(
             chain.append(h)
         return chain
 
+    def _match_prefix(self, seq: Sequence) -> tuple[int, list[int]]:
+        """Read-only usable hit and compressed-prefix hashes, shared with admission.
+
+        Excludes the last block: prefill must still produce sampler logits. The
+        joint SWA/state-checkpoint gate may shorten the compressed hit to a
+        boundary every cache can resume from. Neither sequence bookkeeping nor
+        cache ownership is touched here.
+        """
+        h = -1
+        block_hashes: list[int] = []
+        for i in range(self._n_hash_blocks(seq) - 1):
+            token_ids = self._hash_block_tokens(seq, i)
+            h = self.compute_hash(token_ids, h)
+            block_id = self.kv.lookup(h)
+            if block_id == -1 or self.kv.block(block_id).token_ids != token_ids:
+                break
+            block_hashes.append(h)
+        hit = self._gated_hit(seq, len(block_hashes), block_hashes)
+        return hit, block_hashes
+
+    def prefix_cached_tokens(self, seq: Sequence) -> int:
+        """Reusable prefix tokens, probed without allocating or bookkeeping.
+
+        Only `scheduling_policy="sjf"` calls this, to rank a waiter by the work
+        prefill still has to do instead of by raw prompt length. FCFS never
+        reaches it. Admission probes again through `can_allocate`, because an
+        earlier admission in the same pass can evict the blocks or checkpoints
+        counted here.
+        """
+        if not self.enable_prefix_caching:
+            return 0
+        hit, _ = self._match_prefix(seq)
+        return hit * self.hash_block_size
+
     def can_allocate(self, seq: Sequence, record: bool = True) -> int:
         """Return number of cache-hit blocks (>=0) if seq fits, else -1.
 
@@ -881,18 +915,8 @@ def can_allocate(self, seq: Sequence, record: bool = True) -> int:
             return 0
         # Step 1: compressed prefix (CSA/HCA/indexer share the block hash and
         # read the WHOLE history, so this stays a full front-to-back chained
-        # match). Record each block's hash for the SWA scan below.
-        h = -1
-        compressed_hit = 0
-        block_hashes: list[int] = []
-        for i in range(self._n_hash_blocks(seq) - 1):
-            token_ids = self._hash_block_tokens(seq, i)
-            h = self.compute_hash(token_ids, h)
-            block_id = self.kv.lookup(h)
-            if block_id == -1 or self.kv.block(block_id).token_ids != token_ids:
-                break
-            block_hashes.append(h)
-            compressed_hit += 1
+        # match). `_match_prefix` records each block's hash for the SWA scan.
+        #
         # Step 2: SWA only needs the trailing window before the boundary to be
         # present (SWA is local). Scan right-to-left within the compressed prefix
         # for the largest boundary whose window is SWA-cached (vLLM
@@ -906,7 +930,8 @@ def can_allocate(self, seq: Sequence, record: bool = True) -> int:
         # so a boundary is only resumable where somebody checkpointed the state.
         # `_gated_hit` settles the two gates jointly; neither can be applied to
         # the other's answer.
-        num_cached_blocks = self._gated_hit(seq, compressed_hit, block_hashes)
+        num_cached_blocks, block_hashes = self._match_prefix(seq)
+        compressed_hit = len(block_hashes)
         # Instrumentation: the pre-gate hit, so EngineStats can separate reuse
         # the gates declined (compressed_hit - num_cached_blocks) from reuse
         # lost to compressed eviction (everything above compressed_hit).

atom/model_engine/scheduler.py

@@ -122,6 +122,57 @@ def _optimal_cu_fraction(
         return 0.5
 
 
+SCHEDULING_POLICIES = ("fcfs", "sjf")
+
+
+def validate_scheduling_policy(policy: str) -> str:
+    """Refuse a policy name no scheduler implements.
+
+    Both schedulers select behaviour with `policy == "sjf"`, so an unknown
+    value would otherwise mean FCFS with no diagnostic -- the operator sets the
+    flag, sees the engine start, and gets none of the ordering.
+    """
+    if policy not in SCHEDULING_POLICIES:
+        raise ValueError(
+            f"scheduling_policy must be one of {SCHEDULING_POLICIES}, got {policy!r}"
+        )
+    return policy
+
+
+def shortest_first_key(
+    seq: Sequence, max_skip_steps: int, block_manager: BlockManager | None = None
+) -> tuple[int, int]:
+    """Waiting-queue sort key for `scheduling_policy="sjf"`.
+
+    Two tiers. The front tier holds requests that must not be reordered by
+    length -- a preemption victim, which has already spent a forward pass and
+    had its KV thrown away, and any request skipped `max_skip_steps` times,
+    which is the bound on how long SJF may starve a long prompt. Everything
+    else is ranked by remaining prefill work.
+
+    "Remaining" means after prefix reuse. A fresh waiter owns no blocks and
+    carries `num_cached_tokens == 0`, so `block_manager` is asked what the cache
+    can actually serve it. A waiter that already owns blocks -- an offload resume
+    or a parked partial prefill, and every request in the disaggregated
+    `PrefillScheduler`'s ready set -- instead uses its own computed-token count:
+    that KV is real whether or not it is reachable through the prefix index, so
+    probing would under-report it. Pass `block_manager=None` to skip the probe.
+
+    Both tiers sort stably, so equal-cost requests keep arrival order and the
+    policy degrades to FCFS on a uniform workload.
+
+    Shared by `Scheduler` and `PrefillScheduler`, which are separate classes
+    with separate waiting queues but the same ordering rule.
+    """
+    if seq.is_preempted or 0 < max_skip_steps <= seq.num_skipped_steps:
+        return (0, 0)
+    if block_manager is not None and not seq.block_table:
+        cached = block_manager.prefix_cached_tokens(seq)
+    else:
+        cached = seq.num_cached_tokens
+    return (1, seq.num_tokens - cached)
+
+
 class ScheduledBatch:
     """Immutable snapshot of sequences selected for a single forward pass.
 
@@ -540,6 +591,8 @@ def __init__(
         self._detailed_annotation_enabled = envs.ATOM_ENABLE_DETAILED_ANNOTATION
 
         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
         # Running seqs currently mid-prefill; counter lets schedule() skip the
         # running-queue scan on pure-decode steps.
         self._partial_prefill_count: int = 0
@@ -1127,6 +1180,10 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None:
         num_scheduled_tokens: list[int] = []
         scheduled_spec_decode_tokens: dict[int, np.ndarray] = {}
 
+        # Ordering first, then the two passes that move specific requests to
+        # the front for correctness — their invariants must survive the sort,
+        # so they run after it, not before.
+        self._reorder_waiting_shortest_first()
         self._promote_ready_remote_kv_requests()
         self._park_ready_offload_partial_prefills()
 
@@ -1469,6 +1526,8 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None:
             )
             self.waiting.extend(skipped_waiting_requests)
 
+        self._age_waiting_requests(num_seqs_prefill)
+
         if self._num_parked_remote_kv > 0 and self._schedule_tick % 1000 == 0:
             logger.info(
                 "PD backpressure: parked=%d, waiting=%d, running=%d, "
@@ -2109,6 +2168,9 @@ def _schedule_prefill_seq(
         num_batched_tokens += chunk
         seq.status = SequenceStatus.RUNNING
         seq.type = SequenceType.PREFILL
+        # Admitted: both front-tier claims are spent.
+        seq.is_preempted = False
+        seq.num_skipped_steps = 0
         self.running.append(seq)
         scheduled_seqs[seq.id] = seq
         num_scheduled_tokens.append(chunk)
@@ -2341,6 +2403,9 @@ def preempt(self, seq: Sequence) -> bool:
         # surrendered blocks (the mutator half of `should_defer_free`'s escape).
         self._connector_release_stalled_save(seq)
         self.block_manager.deallocate(seq)
+        # `appendleft` alone no longer guarantees the head: under "sjf" the
+        # next step re-sorts the queue. The flag is what survives that.
+        seq.is_preempted = True
         self.waiting.appendleft(seq)
         return True
 
@@ -3022,6 +3087,50 @@ def _update_waiting_for_remote_kv(self, seq: Sequence) -> bool:
                 )
         return True
 
+    def _reorder_waiting_shortest_first(self) -> None:
+        """Order the waiting queue by remaining prefill work, cheapest first.
+
+        The cost of a request is its uncached token count -- what prefill
+        actually still has to compute, after prefix reuse. A fresh waiter has
+        `num_cached_tokens == 0` until it runs a chunk, so the key would read a
+        cache-hot 20k prompt as 20k of work; `prefix_cached_tokens` probes the
+        live cache read-only instead, under the same SWA/state-checkpoint gate
+        admission uses, so an unresumable prefix is not counted as free.
+
+        Hits are re-probed every sort rather than memoised: an admission earlier
+        in the same pass can evict the blocks a later waiter matched.
+
+        The probe costs one chained-hash walk per fresh waiter per scheduling
+        pass. That is charged to `sjf` alone -- the early return below is what
+        keeps FCFS off this path entirely.
+
+        The sort is stable, so requests of equal cost keep arrival order and
+        the policy degrades to FCFS on a uniform workload.
+        """
+        if self.scheduling_policy != "sjf" or len(self.waiting) < 2:
+            return
+        bm = self.block_manager
+        self.waiting = deque(
+            sorted(
+                self.waiting,
+                key=lambda seq: shortest_first_key(seq, self.sjf_max_skip_steps, bm),
+            )
+        )
+
+    def _age_waiting_requests(self, num_seqs_prefill: int) -> None:
+        """Charge a skip to everything still waiting after an admitting step.
+
+        Gated on a prefill having actually been scheduled: on a step where the
+        engine admitted nobody -- KV exhausted, budget gone, delayer holding --
+        no request was passed over in favour of another, and counting it would
+        let a stall age the whole queue into the front tier at once, which is
+        just FCFS with extra steps.
+        """
+        if self.scheduling_policy != "sjf" or not num_seqs_prefill:
+            return
+        for seq in self.waiting:
+            seq.num_skipped_steps += 1
+
     def _promote_ready_remote_kv_requests(self) -> None:
         """Move completed remote-KV waiters ahead of fresh admissions.
 
@@ -3312,6 +3421,8 @@ def __init__(self, config: Config, disagg_cu_shm_name: str = ""):
         self.block_manager = None  # blocks managed by decode process
         self.waiting: deque[Sequence] = deque()
         self.running: deque[Sequence] = deque()
+        self.scheduling_policy = validate_scheduling_policy(config.scheduling_policy)
+        self.sjf_max_skip_steps = config.sjf_max_skip_steps
         # spec decode not used on prefill side
         self.use_spec = False
         self.spec_decode_local = False
@@ -3399,6 +3510,10 @@ def _schedule(self):
         with self._pending_lock:
             # Collect ready sequences (have received BlockAssignment from decode)
             ready = [s for s in self.waiting if s.block_table]
+            if self.scheduling_policy == "sjf" and len(ready) > 1:
+                ready.sort(
+                    key=lambda seq: shortest_first_key(seq, self.sjf_max_skip_steps)
+                )
 
             for seq in ready:
                 if num_seqs >= self.max_num_seqs:
@@ -3415,6 +3530,13 @@ def _schedule(self):
                 num_batched_tokens += num_new_tokens
                 num_seqs += 1
 
+            # 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:
+                        seq.num_skipped_steps += 1
+
         if not scheduled_seqs:
             return None, {}
 

atom/model_engine/sequence.py

@@ -243,6 +243,17 @@ def __init__(
         # garbage sampled tokens from intermediate chunks and to skip the
         # scheduler's Phase 1 scan when no partials exist.
         self.is_partial_prefill = False
+        # Scheduling steps that admitted some other request's prefill while
+        # this one sat in the waiting queue. Only maintained under
+        # `scheduling_policy="sjf"`, where crossing `sjf_max_skip_steps` moves
+        # the request to the front regardless of length -- the bound on how
+        # long shortest-job-first may starve a long prompt.
+        self.num_skipped_steps = 0
+        # Set by `Scheduler.preempt`, cleared the moment the request is
+        # admitted again. Keeps a preemption victim at the head of the waiting
+        # queue under "sjf", which would otherwise sort it behind every
+        # shorter arrival and re-preempt it forever.
+        self.is_preempted = False
         # `new_block_table` is main's: an array("i") rather than a list,
         # because every forward marshals these into the int32 buffer.
         self.block_table = new_block_table()

docs/configuration_guide.md

@@ -31,6 +31,8 @@ Defined in `atom/config.py`. The root dataclass that the engine consumes.
 | `trust_remote_code` | `bool` | `False` | Trust remote code when loading the model from HuggingFace |
 | `max_num_batched_tokens` | `int` | `16384` | Maximum number of tokens batched together per scheduler step |
 | `scheduler_delay_factor` | `float` | `0.0` | Multiplicative delay (factor x previous prompt latency) before scheduling the next prompt |
+| `scheduling_policy` | `str` | `"fcfs"` | Waiting-queue order: `"fcfs"` (arrival) or `"sjf"` (shortest job first — lowers p90 TTFT, raises it for the longest prompts) |
+| `sjf_max_skip_steps` | `int` | `64` | Under `"sjf"`, promote a request skipped this many steps regardless of length. `0` removes the starvation bound |
 | `max_num_seqs` | `int` | `512` | Maximum number of sequences batched together |
 | `max_model_len` | `int \| None` | `None` | Maximum context length; defaults to `hf_config.max_position_embeddings` (capped by it when set) |
 | `gpu_memory_utilization` | `float` | `0.9` | Fraction of GPU memory available for KV cache and weights (0.0 — 1.0) |