diff --git a/atom/config.py b/atom/config.py index 90ce774f91..9ec52d8cd3 100644 --- a/atom/config.py +++ b/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 diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index f9e1f92164..31feea1402 100644 --- a/atom/model_engine/arg_utils.py +++ b/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, diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index ba3faf1ce1..aba006e7a0 100644 --- a/atom/model_engine/block_manager.py +++ b/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). diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 61a6d2469d..d15921c3c0 100644 --- a/atom/model_engine/scheduler.py +++ b/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, {} diff --git a/atom/model_engine/sequence.py b/atom/model_engine/sequence.py index ffa958bf6d..9c5cf82a89 100644 --- a/atom/model_engine/sequence.py +++ b/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() diff --git a/docs/configuration_guide.md b/docs/configuration_guide.md index 9a452d46c3..6405d22624 100644 --- a/docs/configuration_guide.md +++ b/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) | diff --git a/docs/scheduling_kv_cache_guide.md b/docs/scheduling_kv_cache_guide.md index c9b9cd19dd..880876c5b7 100644 --- a/docs/scheduling_kv_cache_guide.md +++ b/docs/scheduling_kv_cache_guide.md @@ -27,6 +27,8 @@ ATOM (AiTer Optimized Model) uses a prefill-first scheduler with paged KV cache | `kv_cache_block_size` | 16 | Tokens per KV cache block (must be multiple of 16, or 1) | | `enable_prefix_caching` | `False` | Enable hash-based prefix block sharing | | `scheduler_delay_factor` | 0.0 | Delay factor for batching prompt requests (0 = no delay) | +| `scheduling_policy` | `"fcfs"` | Waiting-queue order: `"fcfs"` or `"sjf"` (see below) | +| `sjf_max_skip_steps` | 64 | Starvation bound for `"sjf"` (0 = unbounded) | | `gpu_memory_utilization` | 0.9 | Fraction of GPU memory for KV cache | ## Scheduling algorithm @@ -82,6 +84,63 @@ The scheduler maintains two deques — `waiting` (pending prefill) and `running` 5. Call `block_manager.may_append(seq, num_new_tokens)` where `num_new_tokens = mtp_k + 1`. 6. Re-insert all scheduled sequences back into `running` (preserving order). +### Waiting-queue order (`scheduling_policy`) + +The waiting queue is a FIFO deque. Under the default `"fcfs"` it is left +alone, and requests prefill in arrival order — so one long prompt at the head +adds its whole prefill time to the TTFT of every short request behind it. + +`--scheduling-policy sjf` sorts the queue at the top of `_schedule` by +remaining prefill work, cheapest first +(`Scheduler._reorder_waiting_shortest_first`). Short requests answer sooner +and the longest prompts answer later — the p90 TTFT improves, the p99 gets +worse. Under disaggregated prefill the same ordering runs in +`PrefillScheduler._schedule`, so short requests also reach the decode node's +batch without waiting out a long prompt's prefill. + +Two tiers keep the sort honest (`shortest_first_key`): + +- **Front tier** — preemption victims (which already spent a forward pass and + had their KV freed, so re-sorting them by length would re-preempt them + forever) and any request skipped past `sjf_max_skip_steps`. This is the + starvation bound; `0` removes it, giving pure SJF. +- **Everything else** — ranked by the tokens prefill still has to compute + *after prefix reuse*, stably, so equal-cost requests keep arrival order. + +A skip is only charged on a step that actually admitted a prefill. On a +KV-starved or delayer-held step nobody was passed over, and counting it would +age the whole queue into the front tier at once — FCFS with extra steps. + +**How the cost is measured.** `num_cached_tokens` is advanced by `postprocess`, +so a request that has not yet run a prefill chunk carries 0 — using it alone +would rank a cache-hot 20k prompt as 20k of work and let a cold 500-token +request cut in front. Instead the key splits on whether the waiter owns blocks: + +- **Fresh waiter** (`not seq.block_table`) — `BlockManager.prefix_cached_tokens` + probes the live cache. It shares `_match_prefix` with `can_allocate`, so it + applies the same joint SWA / state-checkpoint gate and never counts a prefix + no cache could resume from. The probe is read-only: no allocation, no + refcount, no `checkpoint_demand_pos`, no funnel counters. +- **Waiter that already owns blocks** — an offload resume, a parked partial + prefill, or anything in the disaggregated `PrefillScheduler`'s ready set uses + its own `num_cached_tokens`. That KV is real whether or not it is reachable + through the prefix index, so probing would under-report it. + +Hits are re-probed on every sort rather than memoised, because an admission +earlier in the same pass can evict blocks a later waiter matched. Admission +probes again for the same reason — the sort's answer is a ranking input, never +an admission decision. + +**Cost, and who pays it.** One chained-hash walk per fresh waiter per +scheduling pass, on top of the walk admission already does. That is charged to +`sjf` alone: `_reorder_waiting_shortest_first` returns before touching the +block manager under any other policy, so FCFS walks the prefix exactly once per +`can_allocate`, as it always has. + +The sort runs *before* `_promote_ready_remote_kv_requests` and +`_park_ready_offload_partial_prefills`, so those passes' front-of-queue +invariants still hold — SJF only orders what they leave alone. + ### Delay factor When `scheduler_delay_factor > 0`, the scheduler delays prefill scheduling to allow the waiting queue to accumulate more requests for better batching: diff --git a/tests/conftest.py b/tests/conftest.py index 7e9c583c21..1dbe0ab4f8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -70,6 +70,8 @@ def __init__(self, **overrides): "eos_token_id": 2, "stop_token_ids": [], "scheduler_delay_factor": 0.0, + "scheduling_policy": "fcfs", + "sjf_max_skip_steps": 64, "speculative_config": None, # Scheduler.__init__ reads config.hf_config.architectures for V4 # SWA-warmup detection; a non-V4 stub keeps that path inert. diff --git a/tests/test_arg_utils_spec.py b/tests/test_arg_utils_spec.py index 82586e9a2b..932d02a391 100644 --- a/tests/test_arg_utils_spec.py +++ b/tests/test_arg_utils_spec.py @@ -5,6 +5,8 @@ import sys from unittest.mock import MagicMock, patch +import pytest + # conftest.py stubs atom.* and zmq before any atom imports are attempted, # but arg_utils.py imports LLMEngine from atom and CompilationConfig / # SpeculativeConfig from atom.config, which the minimal stub doesn't expose. @@ -236,3 +238,29 @@ def test_index_cache_dtype_accepts_underscore_spelling(self): ) assert args.index_cache_dtype == "fp8" + + +class TestSchedulingPolicyCli: + """`--scheduling-policy sjf` must reach the engine kwargs the Config is + built from; a flag that parses but never arrives is the classic silent + no-op for a scheduler knob.""" + + def _parse(self, argv): + parser = argparse.ArgumentParser() + EngineArgs.add_cli_args(parser) + return EngineArgs.from_cli_args(parser.parse_args(argv)) + + def test_default_is_fcfs(self): + assert self._parse([])._get_engine_kwargs()["scheduling_policy"] == "fcfs" + + def test_sjf_reaches_engine_kwargs(self): + kwargs = self._parse(["--scheduling-policy", "sjf"])._get_engine_kwargs() + assert kwargs["scheduling_policy"] == "sjf" + + def test_max_skip_steps_reaches_engine_kwargs(self): + kwargs = self._parse(["--sjf-max-skip-steps", "8"])._get_engine_kwargs() + assert kwargs["sjf_max_skip_steps"] == 8 + + def test_an_unknown_policy_is_refused_at_parse_time(self): + with pytest.raises(SystemExit): + self._parse(["--scheduling-policy", "lifo"]) diff --git a/tests/test_prefill_scheduler.py b/tests/test_prefill_scheduler.py index e1ecf83e64..58a2ceb98e 100644 --- a/tests/test_prefill_scheduler.py +++ b/tests/test_prefill_scheduler.py @@ -219,3 +219,62 @@ def test_schedule_records_prompt_throughput(seq_factory): assert sched.engine_stats.num_prompt_tokens == 4 # Prefill produces no sampled tokens. assert sched.engine_stats.num_generation_tokens == 0 + + +# ── shortest-job-first ordering ─────────────────────────────────────────── + + +def _sjf_scheduler(**overrides): + from atom.model_engine.scheduler import PrefillScheduler + + cfg = {"scheduling_policy": "sjf", "max_num_seqs": 1} + cfg.update(overrides) + return PrefillScheduler(MockConfig(**cfg)) + + +def test_sjf_runs_the_shorter_ready_prefill_first(seq_factory): + """Disaggregated prefill is where ordering pays twice: the short request + finishes sooner AND its decode reaches the decode node's batch sooner, + instead of waiting out a long prompt's prefill.""" + sched = _sjf_scheduler() + long_seq = seq_factory(list(range(40))) + short_seq = seq_factory(list(range(200, 204))) + sched.extend([long_seq, short_seq]) + long_seq.block_table = [0] + short_seq.block_table = [1] + + _, seqs = sched.schedule() + + assert list(seqs) == [short_seq.id] + + +def test_fcfs_remains_the_default_for_prefill_scheduler(seq_factory): + sched = _sjf_scheduler(scheduling_policy="fcfs") + long_seq = seq_factory(list(range(40))) + short_seq = seq_factory(list(range(200, 204))) + sched.extend([long_seq, short_seq]) + long_seq.block_table = [0] + short_seq.block_table = [1] + + _, seqs = sched.schedule() + + assert list(seqs) == [long_seq.id] + + +def test_sjf_promotes_a_repeatedly_skipped_prefill(seq_factory): + """Same starvation bound as the colocated scheduler.""" + sched = _sjf_scheduler(sjf_max_skip_steps=2) + long_seq = seq_factory(list(range(40))) + long_seq.block_table = [0] + sched.add(long_seq) + + order = [] + for i in range(3): + short = seq_factory(list(range(200 + 10 * i, 204 + 10 * i))) + short.block_table = [i + 1] + sched.add(short) + _, seqs = sched.schedule() + order.append(list(seqs)) + + assert order[2] == [long_seq.id] + assert long_seq.id not in order[0] + order[1] diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 5addafff5d..b30a067c28 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -1336,6 +1336,175 @@ def test_partial_prefill_resume_capped(self, seq_factory): assert list(batch2.num_scheduled_tokens) == [8] +# ── shortest-job-first prefill ordering ─────────────────────────────────── + + +class TestSJFScheduling: + """`scheduling_policy="sjf"` orders waiting prefills by remaining work. + + Every scheduler here uses max_num_seqs=1 so exactly one prefill is issued + per step and the batch contents read as an ordering assertion. + """ + + @staticmethod + def _sched(**overrides): + cfg = { + "num_kvcache_blocks": 400, + "kv_cache_block_size": 4, + "max_num_seqs": 1, + "max_num_batched_tokens": 1000, + "max_model_len": 1024, + "scheduling_policy": "sjf", + } + cfg.update(overrides) + return Scheduler(MockConfig(**cfg)) + + def test_short_request_preempts_queue_position_of_earlier_long_request( + self, seq_factory + ): + """The whole point: a 10-token prompt that arrived second prefills + first, so its TTFT is not held hostage by a 100-token prompt.""" + sched = self._sched() + long_seq = seq_factory(list(range(100))) + short_seq = seq_factory(list(range(200, 210))) + sched.add(long_seq) + sched.add(short_seq) + + _, scheduled = sched.schedule() + + assert list(scheduled) == [short_seq.id] + + def test_fcfs_is_the_default_and_keeps_arrival_order(self, seq_factory): + """The policy is opt-in; unset config must schedule the long prompt + first, exactly as before this feature existed.""" + sched = self._sched(scheduling_policy="fcfs") + long_seq = seq_factory(list(range(100))) + short_seq = seq_factory(list(range(200, 210))) + sched.add(long_seq) + sched.add(short_seq) + + _, scheduled = sched.schedule() + + assert list(scheduled) == [long_seq.id] + + def test_equal_length_requests_keep_arrival_order(self, seq_factory): + """The sort is stable, so SJF never reshuffles same-cost requests.""" + sched = self._sched() + first = seq_factory(list(range(10))) + second = seq_factory(list(range(100, 110))) + sched.add(first) + sched.add(second) + + _, scheduled = sched.schedule() + + assert list(scheduled) == [first.id] + + def test_long_request_is_promoted_once_skipped_too_many_times(self, seq_factory): + """SJF alone starves the longest prompt under a stream of short ones. + After `sjf_max_skip_steps` passes it jumps the queue unconditionally, + which is what bounds its worst-case TTFT. + + Budget (not max_num_seqs) throttles this to one prefill per step, so + `running` filling up never confounds the ordering. + """ + sched = self._sched( + max_num_seqs=8, + max_num_batched_tokens=10, + sjf_max_skip_steps=2, + ) + long_seq = seq_factory(list(range(100))) + sched.add(long_seq) + + order = [] + for i in range(3): + sched.add(seq_factory(list(range(200 + 20 * i, 210 + 20 * i)))) + _, scheduled = sched.schedule() + order.append(list(scheduled)) + + # Steps 1 and 2 serve the short arrivals; on step 3 the long request + # has been skipped twice and takes the slot. + assert order[2] == [long_seq.id] + assert long_seq.id not in order[0] + order[1] + + def test_a_step_that_schedules_no_prefill_does_not_age_the_queue(self, seq_factory): + """A KV-starved step passes over nobody — every request is equally + stuck. Counting it as a skip would let a stall promote the whole queue + to the front tier at once, collapsing SJF back into FCFS. + """ + # 25 blocks of 4 tokens: a 100-token prompt is admissible but consumes + # the entire pool, so the next arrival cannot allocate at all. + sched = self._sched( + num_kvcache_blocks=25, max_model_len=128, sjf_max_skip_steps=1 + ) + hog = seq_factory(list(range(100))) + sched.add(hog) + _, scheduled = sched.schedule() + assert list(scheduled) == [hog.id], "setup: the hog must take the pool" + + stalled = seq_factory(list(range(200, 300))) + sched.add(stalled) + _, scheduled = sched.schedule() + + assert scheduled == {}, "setup: no prefill may be scheduled this step" + assert stalled.num_skipped_steps == 0 + + def test_zero_max_skip_steps_disables_the_starvation_bound(self, seq_factory): + """`sjf_max_skip_steps=0` is pure shortest-job-first: no promotion, so + a long prompt yields to short arrivals indefinitely. Documented as a + deliberate setting rather than a degenerate one -- it is the sharpest + p90 win available and the least fair. + """ + sched = self._sched( + max_num_seqs=8, max_num_batched_tokens=10, sjf_max_skip_steps=0 + ) + long_seq = seq_factory(list(range(100))) + sched.add(long_seq) + + served = [] + for i in range(4): + sched.add(seq_factory(list(range(200 + 20 * i, 210 + 20 * i)))) + _, scheduled = sched.schedule() + served.extend(scheduled) + + assert long_seq.id not in served + + def test_preempted_request_is_retried_before_shorter_newcomers(self, seq_factory): + """`preempt` puts its victim at the head on purpose -- it has already + burned a forward pass and its KV was just thrown away. Sorting it back + behind every shorter arrival would re-preempt it indefinitely, so a + preempted request joins the front tier. + """ + sched = self._sched(max_num_seqs=8, max_num_batched_tokens=10) + victim = seq_factory(list(range(100))) + sched.add(victim) + _, scheduled = sched.schedule() + assert list(scheduled) == [victim.id], "setup: victim must run first" + + sched.running.remove(victim) + sched.preempt(victim) + sched.add(seq_factory(list(range(200, 210)))) + + _, scheduled = sched.schedule() + + assert list(scheduled) == [victim.id] + + def test_preemption_flag_clears_once_the_victim_runs_again(self, seq_factory): + """The front tier is a one-shot retry, not a permanent promotion -- + otherwise any request that was ever preempted outranks SJF forever. + """ + sched = self._sched(max_num_seqs=8, max_num_batched_tokens=10) + victim = seq_factory(list(range(100))) + sched.add(victim) + sched.schedule() + sched.running.remove(victim) + sched.preempt(victim) + assert victim.is_preempted is True + + sched.schedule() + + assert victim.is_preempted is False + + # ── prefix caching ──────────────────────────────────────────────────────── @@ -2181,3 +2350,19 @@ def test_no_tier_attached_emits_no_tier_line(self, caplog): with caplog.at_level(logging.INFO, logger="atom"): s._log_pools() assert not any("[Cache Tiers]" in r.getMessage() for r in caplog.records) + + +class TestSchedulingPolicyValidation: + """An unrecognised policy must fail loudly. Both schedulers compare the + string to "sjf", so a typo would otherwise degrade to FCFS in silence -- + the operator sets a flag, sees no error, and gets none of the behaviour.""" + + def test_unknown_policy_is_rejected(self): + with pytest.raises(ValueError, match="scheduling_policy"): + Scheduler(MockConfig(scheduling_policy="SJF")) + + def test_prefill_scheduler_rejects_unknown_policy(self): + from atom.model_engine.scheduler import PrefillScheduler + + with pytest.raises(ValueError, match="scheduling_policy"): + PrefillScheduler(MockConfig(scheduling_policy="shortest")) diff --git a/tests/test_sjf_cache_aware.py b/tests/test_sjf_cache_aware.py new file mode 100644 index 0000000000..ff6fe01ed7 --- /dev/null +++ b/tests/test_sjf_cache_aware.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: MIT +"""SJF ranks actual prefill work, not the full length of a cache-hot prompt.""" + +import pickle + +import pytest +from conftest import MockConfig + +from atom.model_engine.scheduler import ScheduledBatchOutput, Scheduler +from atom.model_engine.sequence import Sequence +from atom.model_engine.state_pool import StateSlotPool +from atom.model_engine.state_runtime import StateTransfer + + +def _scheduler(**overrides): + config = { + "enable_prefix_caching": True, + "kv_cache_block_size": 4, + "num_kvcache_blocks": 40, + "max_num_seqs": 1, + "max_num_batched_tokens": 256, + "max_model_len": 256, + "scheduling_policy": "sjf", + } + config.update(overrides) + return Scheduler(MockConfig(**config)) + + +def _seq(tokens): + return Sequence(tokens, block_size=4) + + +def _warm_cache(sched): + warmup = _seq(list(range(1, 17))) + sched.add(warmup) + batch, scheduled = sched.schedule() + sched.postprocess( + list(scheduled.values()), + ScheduledBatchOutput( + req_ids=[warmup.id], + token_ids=[(2,)], + num_rejected=None, + num_bonus=None, + draft_token_ids=None, + ), + batch=batch, + ) + assert warmup.is_finished + + +def _waiters(sched): + hot = _seq(list(range(1, 17))) # 12 reusable tokens, 4 left to compute. + cold = _seq(list(range(100, 108))) # 8 tokens, all cold. + sched.add(cold) + sched.add(hot) + return hot, cold + + +def test_cache_hot_long_request_beats_cache_cold_short(): + sched = _scheduler() + _warm_cache(sched) + hot, cold = _waiters(sched) + + batch, scheduled = sched.schedule() + + assert list(scheduled) == [hot.id] + assert list(batch.num_scheduled_tokens) == [4] + assert hot.num_cached_tokens == 12 + assert cold.num_cached_tokens == 0 + + +def test_sort_does_not_mutate_sequence_or_cache_bookkeeping(): + sched = _scheduler() + _warm_cache(sched) + hot, cold = _waiters(sched) + before = [pickle.dumps(vars(seq)) for seq in (hot, cold)] + occupancy = sched.block_manager.pool_occupancy() + demand_counters = ( + sched.block_manager.demands_recorded, + sched.block_manager.demands_declined_no_room, + ) + + sched._reorder_waiting_shortest_first() + + assert list(sched.waiting) == [hot, cold] + assert [pickle.dumps(vars(seq)) for seq in (hot, cold)] == before + assert sched.block_manager.pool_occupancy() == occupancy + assert ( + sched.block_manager.demands_recorded, + sched.block_manager.demands_declined_no_room, + ) == demand_counters + + +def test_state_checkpoint_gate_can_reject_an_otherwise_hot_prefix(): + sched = _scheduler() + _warm_cache(sched) + hot, cold = _waiters(sched) + hot.has_per_req_cache = True + state = StateSlotPool(2, StateTransfer.fork(1), hash_block_size=4) + sched.block_manager.state_caches += (state,) + + # KV exists, but there is no matching recurrent state to resume from. + assert sched.block_manager.prefix_cached_tokens(hot) == 0 + sched._reorder_waiting_shortest_first() + assert list(sched.waiting) == [cold, hot] + + # Publish a real checkpoint index entry at the 12-token boundary. + h = -1 + for start in (0, 4, 8): + h = sched.block_manager.compute_hash(hot.token_ids[start : start + 4], h) + state._index(h, 0) + assert sched.block_manager.prefix_cached_tokens(hot) == 12 + sched._reorder_waiting_shortest_first() + assert list(sched.waiting) == [hot, cold] + + +@pytest.mark.parametrize("enable_prefix_caching", [False, True]) +def test_allocated_waiter_uses_computed_tokens(enable_prefix_caching): + """An offload/partial resume owns KV even when it is not prefix-indexed.""" + sched = _scheduler(enable_prefix_caching=enable_prefix_caching) + resume = _seq(list(range(1, 17))) + sched.block_manager.allocate(resume) + resume.num_cached_tokens = 12 + cold = _seq(list(range(100, 108))) + sched.add(cold) + sched.add(resume) + + sched._reorder_waiting_shortest_first() + + assert list(sched.waiting) == [resume, cold] + assert resume.num_cached_tokens == 12 + + +def test_sort_refreshes_after_prefix_eviction(): + sched = _scheduler(num_kvcache_blocks=4) + _warm_cache(sched) + hot, cold = _waiters(sched) + sched._reorder_waiting_shortest_first() + assert list(sched.waiting) == [hot, cold] + + evictor = _seq(list(range(200, 216))) + sched.block_manager.allocate(evictor) + sched._reorder_waiting_shortest_first() + assert list(sched.waiting) == [cold, hot] + + +def test_fcfs_does_not_probe_prefixes(monkeypatch): + sched = _scheduler(scheduling_policy="fcfs") + _warm_cache(sched) + _hot, cold = _waiters(sched) + + def unexpected_probe(seq): + pytest.fail("FCFS must not pay for SJF cache probes") + + monkeypatch.setattr(sched.block_manager, "prefix_cached_tokens", unexpected_probe) + _, scheduled = sched.schedule() + assert list(scheduled) == [cold.id] + + +def _count_prefix_walks(sched, monkeypatch): + """Hash walks per schedule(), counted at the one place they can happen.""" + bm = sched.block_manager + calls = {"match": 0, "admit": 0} + real_match, real_admit = bm._match_prefix, bm.can_allocate + + def counting_match(seq): + calls["match"] += 1 + return real_match(seq) + + def counting_admit(seq): + calls["admit"] += 1 + return real_admit(seq) + + monkeypatch.setattr(bm, "_match_prefix", counting_match) + monkeypatch.setattr(bm, "can_allocate", counting_admit) + return calls + + +def test_fcfs_walks_the_prefix_once_per_admission_attempt(monkeypatch): + """The SJF probe must not become a tax on FCFS. + + `_match_prefix` is the only chained-hash walk in the manager, so counting it + bounds the whole cost. Under FCFS it may run exactly once per `can_allocate` + -- the walk admission has always done -- and never on the sort path. + """ + sched = _scheduler(scheduling_policy="fcfs", max_num_seqs=4) + _warm_cache(sched) + _waiters(sched) + calls = _count_prefix_walks(sched, monkeypatch) + + sched.schedule() + + assert calls["admit"] > 0 + assert calls["match"] == calls["admit"] + + +def test_sjf_pays_one_extra_walk_per_fresh_waiter(monkeypatch): + """The counterpart: SJF's surplus is bounded at one walk per fresh waiter.""" + sched = _scheduler(max_num_seqs=4) + _warm_cache(sched) + waiters = _waiters(sched) + calls = _count_prefix_walks(sched, monkeypatch) + + sched.schedule() + + assert calls["match"] == calls["admit"] + len(waiters) + + +def test_disabled_prefix_caching_preserves_shortest_prompt_order(): + sched = _scheduler(enable_prefix_caching=False) + _warm_cache(sched) + _hot, cold = _waiters(sched) + _, scheduled = sched.schedule() + assert list(scheduled) == [cold.id] + + +@pytest.mark.parametrize("length, cached", [(1, 0), (4, 0), (5, 4), (16, 12)]) +def test_probe_and_admission_leave_the_final_block_to_compute(length, cached): + sched = _scheduler() + _warm_cache(sched) + seq = _seq(list(range(1, length + 1))) + bm = sched.block_manager + assert bm.prefix_cached_tokens(seq) == cached + assert bm.can_allocate(seq) * bm.hash_block_size == cached