reasonsolo · Oct 29, 2026 · Oct 29, 2026 · Oct 30, 2026 · Nov 2, 2026 · Nov 2, 2026
@@ -421,6 +421,35 @@ requests that bypass the handoff retain the configured retry policy. With
 remote-tail replay disabled, ordinary disaggregated requests retain their
 configured retry policy, including the transient TCP retry budget.
 
+### Generation slot preemption (prototype)
+
+GEN slot preemption pauses a resident decode request when all batch slots are
+occupied and a waiter risks missing its TTFT target. Configure exactly one bound
+(neither: disabled):
+
+```yaml
+scheduler_config:
+  gen_preemption:
+    tpot_budget: 0.10          # up to +10% mean TPOT from pauses
+    # target_output_tps: 35.0  # alternative: output-rate floor for ~80% of requests
+```
+
+Preemption requires the auto GEN admission policy
+(`0 < TLLM_GEN_ADMISSION_TTFT_PERCENTILE < 99`) and the V2 KV scheduler.
+Pipeline parallelism and attention DP are rejected. Startup checks these
+prerequisites before reserving paused-request headroom.
+
+Pauses free batch slots while retaining KV, recurrent, and drafter state. Cache
+sizing must cover active plus paused requests: SWA and NVFP4-MLA HP-ring
+reservations include the paused cap; full-attention demand depends on sequence
+lengths. Paused requests are excluded from KV eviction. Disaggregated GEN cannot
+recompute context locally, and Helix CP raises on decode allocation failure.
+Reduce concurrency or increase the KV budget when both slots and pages are full.
+
+Pool sizing assumes one executor per process. Configuring preemption on a context
+worker also reserves headroom. `target_output_tps` has simulation coverage only;
+GPU validation is limited to `tpot_budget`.
+
 ## Environment Variables
 
 TRT-LLM uses some environment variables to control the behavior of disaggregated service.

@@ -153,6 +153,10 @@ def generate_serve_schema() -> dict:
         {**per_layer_dynamic, "required": ["mode"]},
         {"type": ["string", "null"]},
     ]
+    # The model validator allows at most one non-null preemption bound.
+    schema["$defs"]["GenPreemptionConfig"]["anyOf"] = [
+        {"properties": {name: {"type": "null"}}} for name in ("tpot_budget", "target_output_tps")
+    ]
     # These object/Any annotations are Python API hooks, not arbitrary YAML mappings.
     # Update these restrictions when hooks are added or gain YAML construction support.
     for name in ("batched_logits_processor", "checkpoint_loader", "_mpi_session"):

@@ -1098,7 +1098,11 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Opti
     def _get_generation_request_capacity(self) -> int:
         # Pipeline stages retain different in-flight microbatches, and every
         # resident sequence needs one native HP ring page per local layer.
-        return self.max_batch_size * self.mapping.pp_size
+        # GEN slot preemption pauses residents in place, so the paused cap
+        # stays resident too.
+        from tensorrt_llm._torch.pyexecutor.scheduler.gen_preemption import resident_headroom
+
+        return self.max_batch_size * self.mapping.pp_size + resident_headroom(self.max_batch_size)
 
     def get_cache_bytes_per_token(self) -> int:
         sizes, _ = self._get_runtime_cache_size_layer_components()
@@ -1165,7 +1169,13 @@ def get_cache_size_per_token(
             * torch.empty((), dtype=torch.float8_e4m3fn).element_size()
             * num_fp8_layers
         )
-        return cache_bytes, hp_bytes * max_batch_size * mapping.pp_size
+        from tensorrt_llm._torch.pyexecutor.scheduler.gen_preemption import resident_headroom
+
+        # Every resident sequence keeps one native HP ring page per local
+        # layer across pipeline stages; GEN slot preemption pauses residents
+        # in place, so the paused cap stays resident too.
+        resident_capacity = max_batch_size * mapping.pp_size + resident_headroom(max_batch_size)
+        return cache_bytes, hp_bytes * resident_capacity
 
 
 class Fp4MlaKVCacheManagerV2(Fp4MlaV2CacheLayoutPolicy, KVCacheManagerV2):

@@ -997,10 +997,12 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int:
         )
         context_tokens = min(max_tokens, max_context_tokens)
         generation_tokens = max_tokens - context_tokens
+        # Paused decode requests keep their SWA pools; the quota must cover
+        # every resident request, not just the active ones.
         generation_quota = (
             max_tokens * non_sliding_attn_size_per_token
             + generation_tokens * generation_swa_size_per_token
-            + self.max_batch_size * generation_swa_size_per_request
+            + self._get_generation_request_capacity() * generation_swa_size_per_request
         )
         context_extra_quota = context_tokens * context_swa_size_per_token
         padding = self._get_extra_quota_padding()
@@ -1047,7 +1049,9 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:
             use_fp8_ds_mla=self.use_fp8_ds_mla,
         )
         padding = self._get_extra_quota_padding()
-        size_per_batch = self.max_batch_size * generation_swa_size_per_request + padding
+        size_per_batch = (
+            self._get_generation_request_capacity() * generation_swa_size_per_request + padding
+        )
         if quota < size_per_batch:
             return 0
         context_size_per_token = non_sliding_attn_size_per_token + context_swa_size_per_token
@@ -2071,9 +2075,14 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwar
             use_fp8_ds_mla=use_fp8_ds_mla,
         )
         max_batch_size = int(kwargs.get("max_batch_size") or 0)
+        from tensorrt_llm._torch.pyexecutor.scheduler.gen_preemption import resident_headroom
+
+        # Paused decode requests keep their SWA caches; the static cost
+        # covers every resident request, not just the active ones.
+        resident_capacity = max_batch_size + resident_headroom(max_batch_size)
         return (
             non_sliding_attn_size_per_token + swa_size_per_token,
-            swa_size_per_request * max_batch_size,
+            swa_size_per_request * resident_capacity,
         )
 
     def _iter_guard_candidate_buffers(

@@ -89,6 +89,10 @@
                         KVCacheV2Scheduler, MultimodalEagerEncoderScheduler,
                         MultimodalScheduler, SimpleScheduler,
                         SimpleUnifiedScheduler)
+from .scheduler.gen_preemption import GenPreemptionSettings
+from .scheduler.gen_preemption import \
+    resident_headroom as gen_preempt_resident_headroom
+from .scheduler.gen_preemption import unsupported_headroom_reason
 from .seq_slot_manager import SeqSlotManager
 
 GB = 1 << 30
@@ -4236,14 +4240,15 @@ def compute_max_num_sequences(mapping: Mapping,
 
     ``enable_overlap_headroom`` is intentionally opt-in; see
     ``should_enable_overlap_headroom``. Pipeline parallelism already sizes the
-    pool by ``pp_size``.
+    pool by ``pp_size``. GEN slot preemption adds a slot per paused request.
     """
     if mapping.has_pp():
         num_micro_batches = mapping.pp_size
     else:
         num_micro_batches = (2 if enable_overlap_headroom
                              and not disable_overlap_scheduler else 1)
-    return max_batch_size * num_micro_batches
+    return (max_batch_size * num_micro_batches +
+            gen_preempt_resident_headroom(max_batch_size))
 
 
 def resolve_max_num_sequences(model_engine,
@@ -4829,6 +4834,11 @@ def create_py_executor_instance(
         dwdp_manager=dwdp_manager,
         max_num_sequences=max_num_sequences,
     )
+    # Mirror the creator's gen_preemption gate (the pools were sized against
+    # the same reason): an executor that cannot pause keeps the feature off
+    # here instead of re-deriving a short-pool warning from its own check.
+    headroom_reason = unsupported_headroom_reason(
+        llm_args.scheduler_config.gen_preemption, llm_args)
     spec_config = model_engine.spec_config
     config = model_engine.model.model_config.pretrained_config
     return PyExecutor(
@@ -4872,6 +4882,9 @@ def create_py_executor_instance(
         enable_non_overlap_adp_forward_intent=(
             should_enable_non_overlap_adp_forward_intent(
                 mapping, llm_args.disable_overlap_scheduler)),
+        gen_preemption=GenPreemptionSettings.from_config(
+            llm_args.scheduler_config.gen_preemption)
+        if headroom_reason is None else None,
     )
 
 

@@ -2681,6 +2681,9 @@ def _build_inputs(
             total_draft_lens = sum(host.lists.draft_lens)
             spec_metadata.draft_tokens = self._buffers.draft_tokens_cuda[:total_draft_lens]
             spec_metadata.request_ids = host.lists.request_ids
+            spec_metadata.paused_in_place_ids = getattr(
+                scheduled_requests, "paused_in_place_ids", frozenset()
+            )
             spec_metadata.gather_ids = self._gather_ids_buffer.tensor[: len(host.lists.gather_ids)]
             # num_generations / num_tokens / seq_lens are set above, before the
             # attention-DP allgather that must agree with prepare().

@@ -2671,6 +2671,10 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
         self._fresh_fill_announced = False
         self._fresh_fill_unavailable_announced = False
         self.max_admissible_sequences = max_num_sequences * (2 if needs_extra_leases else 1)
+        # Requests paused in place by GEN slot preemption stay resident.
+        from ..scheduler.gen_preemption import resident_headroom
+
+        self.max_admissible_sequences += resident_headroom(max_batch_size)
         # The guard page is held by a permanent sequence, so it needs an index
         # slot of its own. Taking one of the scheduler's would change which
         # requests get admitted, so a run with the diagnostic on would no
@@ -3328,7 +3332,9 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Opti
 
     def _get_generation_request_capacity(self) -> int:
         """Return the resident generation requests used for SWA sizing."""
-        return self.max_batch_size
+        from ..scheduler.gen_preemption import resident_headroom
+
+        return self.max_batch_size + resident_headroom(self.max_batch_size)
 
     def _get_max_tokens_from_quota(self, quota: int) -> float:
         """Rank-local byte quota -> token capacity (GLOBAL tokens under helix)."""
@@ -7839,8 +7845,11 @@ def get_cache_size_per_token(
         )
         # The affine slope covers all tokens; context additionally retains SWA
         # pages for the current token batch beyond the generation windows.
+        from ..scheduler.gen_preemption import resident_headroom
+
+        resident_capacity = max_batch_size + resident_headroom(max_batch_size)
         fixed_cost = (
-            swa_size_per_request * max_batch_size
+            swa_size_per_request * resident_capacity
             + (context_size_per_token - cache_size_per_token) * max_num_tokens
         )
         bytes_per_slot = _get_single_swa_pool_slot_bytes(
@@ -7866,7 +7875,7 @@ def get_cache_size_per_token(
                     )
                     costs.append(
                         (
-                            layer_size * max_num_tokens + retained_bytes * max_batch_size,
+                            layer_size * max_num_tokens + retained_bytes * resident_capacity,
                             layer_size * tokens_per_block,
                         )
                     )

@@ -635,7 +635,9 @@ def _estimate_mamba_hybrid_cache_cost(
         )
     state_bytes_per_rank = local_mamba_layers * params.get_states_bytes_per_layer(mapping)
     state_bytes_per_rank += extra_state_bytes_per_rank
-    max_resident_sequences = max_batch_size * mapping.pp_size
+    from ...scheduler.gen_preemption import resident_headroom
+
+    max_resident_sequences = max_batch_size * mapping.pp_size + resident_headroom(max_batch_size)
 
     if include_explicit_snapshots:
         fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts(

@@ -638,7 +638,10 @@ def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]:
         return super()._get_pool_roles(pool_id)
 
     def _max_resident_sequences(self) -> int:
-        return self.max_batch_size * self.mapping.pp_size
+        from ...scheduler.gen_preemption import resident_headroom
+
+        # Plus the requests paused in place by GEN slot preemption.
+        return self.max_batch_size * self.mapping.pp_size + resident_headroom(self.max_batch_size)
 
     def _mamba_state_bytes_per_slot(self) -> int:
         """Return persistent bytes per slot allocated by the native pools."""

@@ -133,6 +133,8 @@
                                    count_retiring_requests)
 from .scheduler.gen_admission_queue import (GenAdmissionConfig,
                                             GenAdmissionWaitingQueue)
+from .scheduler.gen_preemption import (GenPreemptionSettings,
+                                       create_gen_preemption)
 from .scheduler.scheduler import RemoteTailPhase
 from .scheduler.scheduler_v2 import KVCacheV2Scheduler
 
@@ -700,6 +702,7 @@ def __init__(
         enable_adp_dummy_fixes: bool = False,
         enable_scheduler_aware_adp_dummy: bool = False,
         enable_non_overlap_adp_forward_intent: bool = False,
+        gen_preemption: Optional[GenPreemptionSettings] = None,
         transfer_budget_is_explicit: bool = False,
     ) -> None:
         if kv_pool_rebalance_check_interval <= 0:
@@ -800,6 +803,7 @@ def __init__(
         )
         self._resources_released = False
         self._closed = False
+        self._gen_preemption_settings = gen_preemption
         self._initialize_execution(
             PyExecutorResources(
                 resource_manager=resource_manager,
@@ -1411,6 +1415,9 @@ def on_detected():
         if self.dist.rank == 0 and isinstance(self.waiting_queue,
                                               GenAdmissionWaitingQueue):
             self._first_token_timer = self.waiting_queue.enable_rank0()
+        # GEN slot preemption (scheduler_config.gen_preemption).
+        self._gen_preemption = create_gen_preemption(
+            self, getattr(self, "_gen_preemption_settings", None))
 
         self.control_request_barrier = threading.Event()
         self.control_action_done = threading.Event()
@@ -8209,6 +8216,19 @@ def _fetch_new_requests(
         self._fetch_and_enqueue_requests(waiting_queue, total_num_live_requests)
         self._cancel_queued_context_requests(waiting_queue)
 
+        # 2b. GEN slot preemption: resume paused requests, then pause running
+        # ones for at-risk waiting requests.
+        preemption = getattr(self, "_gen_preemption", None)
+        if preemption is not None:
+            total_num_active_requests = preemption.step(
+                active_requests, self.request_broadcaster.rank0_now,
+                self._get_request_admission_capacity(),
+                total_num_active_requests, waiting_queue)
+            self.scheduler.set_gen_paused_ids(preemption.paused_ids,
+                                              preemption.resumed_ids)
+            if self.dist.rank == 0:
+                preemption.maybe_log()
+
         # 3. Pop requests from waiting queue
         new_requests, placement = self._pop_from_waiting_queue(
             waiting_queue, total_num_active_requests, all_rank_states)
@@ -8772,8 +8792,12 @@ def _schedule(self):
             self.kv_cache_manager.prepare_expect_snapshot_points(
                 self.active_requests)
 
+        requests_to_schedule = self.active_requests
+        if getattr(self, "_gen_preemption", None) is not None:
+            requests_to_schedule = self._gen_preemption.scheduling_order(
+                requests_to_schedule)
         scheduler_output = self.scheduler.schedule_request(
-            self.active_requests, self.inflight_req_ids)
+            requests_to_schedule, self.inflight_req_ids)
 
         scheduled_encoder_requests = scheduler_output.encoder_requests
         should_batch_encoder_requests = (self.is_encoder_decoder
@@ -8830,6 +8854,27 @@ def _schedule(self):
         scheduled_requests.scheduled_mm_encoder_items = (
             scheduler_output.scheduled_mm_encoder_items)
         scheduled_requests.recompute_paused_requests = scheduler_output.recompute_paused_requests
+        if getattr(self, "_gen_preemption", None) is not None:
+            retained_ids = set(self._gen_preemption.paused_ids)
+            scheduled_ids = {
+                req.py_request_id
+                for req in scheduled_requests.generation_requests
+            }
+            evicted_ids = {
+                req.py_request_id
+                for req in (scheduled_requests.paused_requests +
+                            scheduled_requests.recompute_paused_requests)
+            }
+            for req in self.active_requests:
+                if (req.state == LlmRequestState.GENERATION_IN_PROGRESS
+                        and req.py_request_id not in scheduled_ids
+                        and req.py_request_id not in evicted_ids):
+                    # Forced resumes may displace another live decode request.
+                    # Its previous sample is harvested without a new forward.
+                    req.py_batch_idx = None
+                    self._gen_preemption.record_displaced(req)
+                    retained_ids.add(req.py_request_id)
+            scheduled_requests.paused_in_place_ids = frozenset(retained_ids)
 
         self._maybe_record_hang_diagnostic_phase("scheduled",
                                                  scheduled_requests)

@@ -62,6 +62,7 @@
 from .model_engine import ModelEngine
 from .model_loader import ModelLoader, _construct_checkpoint_loader
 from .py_executor import PyExecutor
+from .scheduler import gen_preemption
 
 _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS = (90, 100, 103, 107, 120, 121)
 _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS = (90, 100, 103, 107, 120)
@@ -604,6 +605,14 @@ def _create_py_executor(
     kv_connector_config = llm_args.kv_connector_config
 
     scheduler_config = llm_args.scheduler_config
+    # Reserve paused-request headroom only when startup prerequisites allow it.
+    headroom_reason = gen_preemption.unsupported_headroom_reason(
+        scheduler_config.gen_preemption, llm_args)
+    if headroom_reason:
+        logger.warning(
+            f"[gen_preempt] GEN slot preemption disabled: {headroom_reason}")
+    gen_preemption.configure(
+        None if headroom_reason else scheduler_config.gen_preemption)
 
     # Since peft_cache_config may be subject to change, avoid these changes propagate back
     # to llm_args.peft_cache_config

@@ -420,6 +420,27 @@ def _next(self) -> tuple:
             self._order = self._build_order()[::-1]
         return self._order[-1]
 
+    @property
+    def target_s(self) -> Optional[float]:
+        """The ``auto`` policy's TTFT target X once it is set, else None."""
+        return self.adaptive.X if self.adaptive is not None else None
+
+    def num_at_risk(self, slack_s: float) -> int:
+        """Waiting requests with at most ``slack_s`` to spare against X (late
+        ones included) if the next pop is one of them, else 0 (GEN slot
+        preemption)."""
+        ada, now = self.adaptive, self._now
+        if ada is None or ada.X is None or now is None or not self:
+            return 0
+
+        def at_risk(entry) -> bool:
+            orig, _, _, keyed, _ = entry
+            return keyed and ada.X - (now - orig + ada.s2f_s) <= slack_s
+
+        if not at_risk(self._next()[0]):
+            return 0
+        return sum(1 for e in self._entries if e[1] not in self._popped and at_risk(e))
+
     def mark_admitted(self, items: Iterable) -> None:
         now = self._now
         for item in items: