feat(scheduler): interleave decode steps into long prefill runs

python/freetoken/engine/config.py

@@ -77,6 +77,11 @@ class EngineConfig:
     # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime
     # rebuild can pin an absolute window instead).
     swa_full_tokens_ratio: float = 0.2
+    # Force one decode step after this many consecutive prefill steps, so a long chunked
+    # prefill cannot starve in-flight decodes. None/0 keeps the historical prefill-first
+    # order (the scheduler's own TODO names this: "support other policies: e.g. DECODE
+    # first"). At 8 the cost is ~1% of prefill wall time.
+    decode_interleave_every: int | None = None
     # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.

python/freetoken/scheduler/interleave.py

@@ -0,0 +1,73 @@
+"""Decode-interleave scheduling policy.
+
+Background (measured on production, 2026-09-15)
+-----------------------------------------------
+``Scheduler._schedule_next_batch`` picks::
+
+    batch = (prefill_manager.schedule_next_batch(budget)
+             or decode_manager.schedule_next_batch())
+
+Prefill wins unconditionally.  A long prompt is chunked at the window pool's
+budget (384 tokens in production), so a 300k-token context is 797 consecutive
+prefill steps -- and while they run, *decoding requests already in flight are
+never scheduled*.  Measured on the live service:
+
+    four long cold prefills: 680/797/506/250 chunks, 267/317/190/93 s
+    prefill-burst share of a 15-minute window: 70%
+
+so an in-flight request waits up to ~5 minutes for its next token.  The
+scheduler already carries the marker for this::
+
+    # TODO: support other policies: e.g. DECODE first
+
+Reversing the order outright is wrong -- it starves prefill, and a request that
+never prefills never starts.  What is needed is *interleaving*: let prefill keep
+priority, but force a decode step every N prefill steps so in-flight requests
+make progress.  The cost is small by construction: measured chunk cost is
+0.388 s, and one decode step for a handful of requests is a fraction of that,
+so every 8th step costs ~1% of prefill wall time while bounding the stall at
+8 chunks (~3 s) instead of the whole burst (~300 s).
+
+This module holds the *policy* only -- a counter and a decision -- so it can be
+unit-tested without a Scheduler, an engine, or a GPU.
+"""
+
+from __future__ import annotations
+
+
+class DecodeInterleavePolicy:
+    """Decide whether to force a decode step between prefill steps.
+
+    ``every`` is the number of consecutive prefill steps allowed before a decode
+    step is forced.  ``None`` or <= 0 disables the policy, reproducing the
+    historical prefill-first behaviour bit-for-bit.
+
+    The counter only advances on prefill steps that were actually scheduled, so
+    a scheduler with nothing to prefill never accumulates credit, and the first
+    decode of a burst is not delayed.
+    """
+
+    def __init__(self, every: int | None = None) -> None:
+        self.every = int(every) if every else None
+        self._prefill_streak = 0
+
+    @property
+    def enabled(self) -> bool:
+        return self.every is not None and self.every > 0
+
+    def note_prefill(self) -> None:
+        """Record that a prefill step was scheduled in this slot."""
+        self._prefill_streak += 1
+
+    def note_decode(self) -> None:
+        """Record that a decode step was scheduled in this slot (streak resets)."""
+        self._prefill_streak = 0
+
+    def wants_decode(self) -> bool:
+        """True when the streak has reached the threshold and a decode is due."""
+        if not self.enabled:
+            return False
+        return self._prefill_streak >= self.every
+
+    def reset(self) -> None:
+        self._prefill_streak = 0

python/freetoken/scheduler/scheduler.py

@@ -31,6 +31,7 @@
 from .cache import CacheManager
 from .config import SchedulerConfig
 from .decode import DecodeManager
+from .interleave import DecodeInterleavePolicy
 from .io import SchedulerIOMixin
 from .mm import cut_image_spans, plan_mm_batch
 from .prefill import ChunkedReq, PrefillManager
@@ -91,6 +92,14 @@ def __init__(self, config: SchedulerConfig):
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
         )
         self.decode_manager = DecodeManager(config.page_size)
+        # Bound how long a run of prefill steps may starve in-flight decodes. A long
+        # prompt chunked at the window budget is hundreds of consecutive prefill steps;
+        # measured on an 8xRTX4090 DSV4 deployment, that left an already-decoding
+        # request unscheduled for 267-317 s. Unset reproduces the historical
+        # prefill-first order exactly.
+        self._interleave = DecodeInterleavePolicy(
+            getattr(config, "decode_interleave_every", None)
+        )
         self._bidirectional_mm = any(getattr(g, "bidirectional_mm_blocks", False) for g in config.model_config.attention_groups)
         self.prefill_manager = PrefillManager(
             self.cache_manager,
@@ -867,11 +876,28 @@ def _gather_multimodal(self, batch: Batch) -> None:
             )
 
     def _schedule_next_batch(self) -> ForwardInput | None:
-        # TODO: support other policies: e.g. DECODE first
-        batch = (
-            self.prefill_manager.schedule_next_batch(self.prefill_budget)
-            or self.decode_manager.schedule_next_batch()
-        )
+        # Prefill keeps priority -- a request that never prefills never starts -- but a
+        # long run of prefill chunks must not starve in-flight decodes. When the policy
+        # says a decode is due, take one if a decode is actually runnable; otherwise stay
+        # with prefill (never spend a slot idle just to keep a promise).
+        # getattr, not attribute access: the accounting tests drive this method on a
+        # partially built Scheduler, and interleaving must be an opt-in that a stripped
+        # object simply does not have rather than something that breaks it. A missing
+        # policy is the historical prefill-first order.
+        policy = getattr(self, "_interleave", None)
+        batch = None
+        if policy is not None and policy.wants_decode():
+            batch = self.decode_manager.schedule_next_batch()
+            if batch is not None:
+                policy.note_decode()
+        if batch is None:
+            batch = self.prefill_manager.schedule_next_batch(self.prefill_budget)
+            if batch is not None and policy is not None:
+                policy.note_prefill()
+        if batch is None:
+            batch = self.decode_manager.schedule_next_batch()
+            if batch is not None and policy is not None:
+                policy.note_decode()
         if batch is None:
             return None
         forward_input = self._prepare_batch(batch)

python/freetoken/server/args.py

@@ -407,6 +407,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         ),
     )
 
+    parser.add_argument(
+        "--decode-interleave-every",
+        type=int,
+        default=ServerArgs.decode_interleave_every,
+        help=(
+            "Force one decode step after this many consecutive prefill steps. A long "
+            "prompt is prefilled in chunks, so it occupies that many consecutive "
+            "scheduler steps; measured on an 8xRTX4090 DSV4 deployment, a 300k-token "
+            "context left an already-decoding request unscheduled for 267-317 s. "
+            "Unset keeps the historical prefill-first order (the scheduler's own TODO "
+            "names this policy)."
+        ),
+    )
+
     parser.add_argument(
         "--page-size",
         type=int,