perf(scheduler): coalesce prefills on TP

atom/model_engine/block_manager.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from math import inf, isinf
 from time import monotonic
+from weakref import WeakKeyDictionary
 
 import numpy as np
 import xxhash
@@ -131,6 +132,11 @@ def __init__(
         # tokens (see _hash_block_size). == block_size when DCP is off.
         self.hash_block_size = self.block_size * self.dcp_world_size
         self.enable_prefix_caching = config.enable_prefix_caching
+        # Content hashes only: pool hits and resource fit are always rechecked.
+        # Weak keys keep this scheduler-side cache out of request serialization.
+        self._prefill_probe_hashes: WeakKeyDictionary[
+            Sequence, tuple[int, list[int]]
+        ] = WeakKeyDictionary()
         self.total_evicted_blocks: int = 0
 
         kv_events = getattr(config, "kv_events_config", None)
@@ -861,17 +867,22 @@ def _chain_to(
             chain.append(h)
         return chain
 
-    def can_allocate(self, seq: Sequence, record: bool = True) -> int:
+    def can_allocate(
+        self,
+        seq: Sequence,
+        record: bool = True,
+        *,
+        block_hashes: list[int] | None = None,
+        reuse_hashes: bool = False,
+    ) -> int:
         """Return number of cache-hit blocks (>=0) if seq fits, else -1.
 
-        `record=False` marks a fit probe -- asking only whether the seq *could*
-        be admitted. The fit answer is identical, but the probe is not
-        side-effect-free: the instrumentation and checkpoint-demand/-end writes
-        below run before the `record` gate. That is safe -- the sole probe caller
-        reads only the `>= 0` return, and the next real `can_allocate` overwrites
-        those fields first. `record` gates only the joint-boundary commit (seq
-        joint fields + funnel counters), so a probe cannot inflate the
-        operator-visible funnel (see `_commit_joint_boundary`).
+        `record=False` returns the same fit and HBM hit without committing
+        joint-load fields or joint-boundary counters. Checkpoint demand/end
+        and hit instrumentation still refresh; demand counters deduplicate per
+        request. An optional `block_hashes` output
+        lets the scheduler defer `record_allocation` until after its wait and
+        token-budget checks, reusing this probe's chain without another walk.
 
         The hit count is the contiguous run of cache hits starting at the
         prompt's first block. On the first miss we break: subsequent blocks
@@ -888,6 +899,10 @@ def can_allocate(self, seq: Sequence, record: bool = True) -> int:
         # The full per-request width, because that is what `allocate` will take:
         # gating on one slot would admit a request the pool cannot give a
         # rollback set to.
+        if block_hashes is None:
+            block_hashes = []
+        else:
+            block_hashes.clear()
         if seq.has_per_req_cache and not self.state.has_free(self.state_slots_per_req):
             return -1
         if not self.enable_prefix_caching:
@@ -899,10 +914,21 @@ def can_allocate(self, seq: Sequence, record: bool = True) -> int:
         # match). Record each block's hash for the SWA scan below.
         h = seq.cache_seed
         compressed_hit = 0
-        block_hashes: list[int] = []
+        cached_hashes = None
+        if reuse_hashes:
+            seed, cached_hashes = self._prefill_probe_hashes.get(seq, (h, []))
+            if seed != h:
+                cached_hashes = []
+            self._prefill_probe_hashes[seq] = (h, cached_hashes)
+        immutable_blocks = seq.num_prompt_tokens // self.hash_block_size
         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)
+            if cached_hashes is not None and i < len(cached_hashes):
+                h = cached_hashes[i]
+            else:
+                h = self.compute_hash(token_ids, h)
+                if cached_hashes is not None and i < immutable_blocks:
+                    cached_hashes.append(h)
             block_id = self.kv.lookup(h)
             if block_id == -1 or self.kv.block(block_id).token_ids != token_ids:
                 break
@@ -949,25 +975,11 @@ def can_allocate(self, seq: Sequence, record: bool = True) -> int:
         self._record_checkpoint_end(seq)
         if not self._has_page_units(num_new_blocks, protected_hash):
             return -1
-        # A boundary LMCache and the tier can jointly reach, above this hit.
-        # Committed to the seq rather than returned: what `allocate` claims from
-        # HBM is still `num_cached_blocks`, and the joint boundary only decides
-        # where the two loads are aimed. `record` is the probe/admission split:
-        # a real admission computes the boundary and commits it (seq fields +
-        # funnel counters); a fit probe skips it entirely. The `num_cached_blocks`
-        # returned below does not depend on the decision, so gating the
-        # computation -- not just the commit -- spares a probe the O(prompt) work
-        # for a value it would only discard: `_joint_kv_boundary` walks a chained
-        # xxhash up to the LMCache-only cap (`_chain_to`) plus a `_gated_hit`
-        # rescan. A 128k prompt at the front of a KV-pressured queue paid that
-        # full chain on every scheduling pass (`is_mixed_batch` peeks up to four
-        # waiting seqs with `record=False`) for a decision nothing read. Placed
-        # below the refusal, not above it, for the same reason `_extend_hash_chain`
-        # is: a refused admission would discard it, and nothing between here and
-        # the refusal reads the seq's joint fields -- this is that move's twin.
+        # A direct admission commits its joint boundary here. The scheduler
+        # probes first and calls record_allocation only after dependency and
+        # budget checks; refused probes avoid the extra LMCache hash walk.
         if record:
-            decision = self._joint_kv_boundary(seq, num_cached_blocks, block_hashes)
-            self._commit_joint_boundary(seq, decision)
+            self.record_allocation(seq, num_cached_blocks, block_hashes)
         # After the refusal, not before it. The chain is O(prompt) xxhash plus
         # two temporaries per block, and a refused admission discards it — a
         # 128k prompt queued behind a full pool paid ~2000 rounds per waiting
@@ -983,6 +995,14 @@ def can_allocate(self, seq: Sequence, record: bool = True) -> int:
         self._extend_hash_chain(seq, block_hashes)
         return num_cached_blocks
 
+    def record_allocation(
+        self, seq: Sequence, num_cached_blocks: int, block_hashes: list[int]
+    ) -> None:
+        """Commit the fit probe before any allocation changes the pool."""
+        if self.enable_prefix_caching:
+            decision = self._joint_kv_boundary(seq, num_cached_blocks, block_hashes)
+            self._commit_joint_boundary(seq, decision)
+
     def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool:
         """Allocate blocks for `seq`. `num_cached_blocks` is the hit count
         returned by `can_allocate` (0 if caller didn't call it).
@@ -2060,7 +2080,9 @@ def cancel_midstep(self, seq: Sequence) -> None:
         self.state.cancel_midstep(seq.midstep_reservations)
         seq.midstep_reservations = []
 
-    def checkpoint_cut(self, seq: Sequence, start: int, end: int) -> int:
+    def checkpoint_cut(
+        self, seq: Sequence, start: int, end: int, *, record: bool = True
+    ) -> int:
         """Earliest ladder position in `(start, end]`, or 0 if there is none.
 
         What a prefill chunk is cut at so its forward lands exactly on a rung.
@@ -2133,7 +2155,7 @@ class still needs the forward to land there, and the readable ones lose
         # `chunks_cut_for_demand` would swamp the convergence signal that
         # counter exists to expose. The demand is checked first because when
         # the two coincide it is the demand that evidenced the position.
-        if target != rung and target < end:
+        if record and target != rung and target < end:
             if target == demand:
                 self.chunks_cut_for_demand += 1
             else:
@@ -2607,6 +2629,7 @@ def deallocate(self, seq: Sequence):
             seq.state_fork_src = -1
 
         seq.offload_joint.load_hash = -1
+        seq.offload_joint.reset_joint()
 
     def can_append(self, seq: Sequence, num_new_tokens: int = 1) -> bool:
         seq_len = len(seq)

atom/model_engine/engine_core.py

@@ -167,6 +167,12 @@ def __init__(self, config: Config, input_address: str, output_address: str):
                 config,
                 state_runtime=self.state_runtime,
             )
+            if (
+                config.parallel_config.data_parallel_size == 1
+                and config.pipeline_parallel_size == 1
+                and envs.ATOM_PREFILL_DECODE_INTERVAL > 0
+            ):
+                self._init_prefill_delayer(config)
 
         self.kv_transfer_enabled = bool(config.kv_transfer_config)
         self._next_idle_kv_drain = 0.0
@@ -188,6 +194,31 @@ def __init__(self, config: Config, input_address: str, output_address: str):
         self._send_ready_signal()
         logger.info(f"{self.label}: EngineCore fully initialized and ready")
 
+    def _init_prefill_delayer(self, config: Config, cpu_group=None):
+        if (
+            not envs.ATOM_ENABLE_PREFILL_DELAYER
+            or config.enable_rapidserve
+            or self.scheduler is None
+        ):
+            return
+        from atom.model_engine.prefill_delayer import PrefillDelayer
+
+        self.scheduler.set_prefill_delayer(
+            PrefillDelayer(
+                dp_size=config.parallel_config.data_parallel_size,
+                cpu_group=cpu_group,
+                max_num_batched_tokens=config.max_num_batched_tokens,
+                target_fill=envs.ATOM_PREFILL_DELAYER_TARGET_FILL,
+                ttft_max_ticks=envs.ATOM_PREFILL_DELAYER_TTFT_MAX_TICKS,
+                partial_max_ticks=envs.ATOM_PREFILL_DELAYER_PARTIAL_MAX_TICKS,
+                stall_ticks=envs.ATOM_PREFILL_DELAYER_STALL_TICKS,
+                kv_high_watermark=envs.ATOM_PREFILL_DELAYER_KV_HIGH_WATERMARK,
+                token_usage_low_watermark=envs.ATOM_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK,
+                max_queue_ms=envs.ATOM_PREFILL_DELAYER_MAX_QUEUE_MS,
+                prefill_decode_interval=envs.ATOM_PREFILL_DECODE_INTERVAL,
+            )
+        )
+
     def _freeze_after_startup(self):
         """Freeze this process and its ModelRunner workers.
 
@@ -726,24 +757,7 @@ def __init__(self, config: Config, input_address: str, output_address: str):
         self.engines_running = True
         self._shutting_down = False
 
-        if envs.ATOM_ENABLE_PREFILL_DELAYER:
-            from atom.model_engine.prefill_delayer import PrefillDelayer
-
-            self.scheduler.set_prefill_delayer(
-                PrefillDelayer(
-                    dp_size=config.parallel_config.data_parallel_size,
-                    cpu_group=self.dp_group,
-                    max_num_batched_tokens=config.max_num_batched_tokens,
-                    target_fill=envs.ATOM_PREFILL_DELAYER_TARGET_FILL,
-                    ttft_max_ticks=envs.ATOM_PREFILL_DELAYER_TTFT_MAX_TICKS,
-                    partial_max_ticks=envs.ATOM_PREFILL_DELAYER_PARTIAL_MAX_TICKS,
-                    stall_ticks=envs.ATOM_PREFILL_DELAYER_STALL_TICKS,
-                    kv_high_watermark=envs.ATOM_PREFILL_DELAYER_KV_HIGH_WATERMARK,
-                    token_usage_low_watermark=envs.ATOM_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK,
-                    max_queue_ms=envs.ATOM_PREFILL_DELAYER_MAX_QUEUE_MS,
-                    prefill_decode_interval=envs.ATOM_PREFILL_DECODE_INTERVAL,
-                )
-            )
+        self._init_prefill_delayer(config, self.dp_group)
 
     def _init_data_parallel(self, config: Config):
         dp_rank = config.parallel_config.data_parallel_rank

atom/model_engine/engine_utility.py

@@ -218,16 +218,19 @@ def _handle_clear_kv_cache(self, args: dict):
         )
 
     def _handle_abort_request(self, args: dict):
-        """Mark a sequence ABORTED (client disconnected) so the scheduler finishes
-        it at the next step via the normal stop path (frees KV, drops it)."""
+        """Cancel queued work promptly; running forwards finish via postprocess."""
         req_id = args.get("req_id") if isinstance(args, dict) else None
         if req_id is None or self.scheduler is None:
             return
-        found = False
-        for seq in list(self.scheduler.running) + list(self.scheduler.waiting):
-            if seq.id == req_id:
-                seq.status = SequenceStatus.ABORTED
-                found = True
+        abort = getattr(self.scheduler, "abort_request", None)
+        if callable(abort):
+            found = abort(req_id)
+        else:
+            found = False
+            for seq in list(self.scheduler.running) + list(self.scheduler.waiting):
+                if seq.id == req_id:
+                    seq.status = SequenceStatus.ABORTED
+                    found = True
         logger.info(f"{self.label}: abort_request req_id={req_id} found={found}")
 
     def _handle_configure_hidden_states(self, args: dict):