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

[huhoo]

Closes #483.

Summary

Scheduler._schedule_next_batch gives prefill unconditional priority via an or short-circuit, so decode is never consulted while a long prompt's chunked prefill is running. On an 8x RTX 4090 DSV4 deployment, a request already decoding went 77 s without being scheduled at all while one 200K-token prompt prefilled.

This adds --decode-interleave-every N: prefill keeps priority, but after N consecutive prefill steps one decode step is taken if a decode is runnable. Measured on the production box with the flag enabled, the worst inter-token gap went from a single 77 s stall to 13.03 s.

This is the policy the scheduler's own TODO names (support other policies: e.g. DECODE first). Reversing the priority outright is not the fix — it would starve prefill, and a request that never prefills never starts.

Design

• The counter advances only 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.
• When the policy says a decode is due but no decode is runnable, prefill proceeds and the streak is not reset. A slot is never spent idle to keep the promise.
• The policy is read through getattr because the accounting tests drive _schedule_next_batch on a Scheduler built with __new__. A stripped object keeps the historical prefill-first order instead of raising.
• DecodeInterleavePolicy is a pure counter + decision in its own module, so it is testable without a Scheduler, an engine, or a GPU.

Measured on hardware

The flag has now been deployed on the production box. Controlled experiment, temperature: 0, no other traffic, current production geometry (--swa-full-tokens-ratio 0.12 -> 308 window pages -> 3456-token chunk budget, --max-running-requests 50):

• victim — short prompt, streams ~700 tokens. Measured as the gap between consecutive streamed chunks.
• attacker — one 200,019-token prompt, fired 6 s after the victim, asking for 3 tokens.

	attacker	victim	max inter-token gap	gaps > 5 s	gaps > 20 s
victim alone (baseline)	—	11.7 s / 699 tok	0.03 s	0	0
flag unset	79.5 s	90.4 s / 699 tok	77 s (see below)	1	1
flag = 8	79.7 s	90.6 s / 699 tok	13.03 s	7	0

Three independent checks agree that this is the mechanism and not noise:

1. The attacker's 58 chunks at N=8 should interleave 58 / 8 = 7.25 decode steps. Measured gaps > 5 s = 7.
2. The bound predicted from the measured per-chunk cost (8 chunks x 1.35 s = 10.8 s) against the observed 13.03 s.
3. Total wall time is unchanged (90.4 s -> 90.6 s) by design — this is not a throughput change. The victim's ~11 s of decode and the attacker's ~79 s of prefill are the same work; interleaving only moves the decode from "all after the prefill" to "spread through it". The gain is responsiveness: the longest wait for the next token drops from ~77 s to ~13 s.

For the unset row the gap is established from the engine log rather than a streamed run: 58 consecutive prefill steps over 77 s with the victim queued throughout, and no decode step in that span.

The earlier numbers in this PR were wrong and I have corrected them (see #483): the 267 s / 384-token-chunk figures were from an older --swa-full-tokens-ratio 0.02 configuration. The corrected expectation for this flag is 77 s -> ~11 s, not 267 s -> 3 s.

Testing

• tests/scheduler/test_interleave.py (30 tests).
• The truth table drives a real Scheduler via __new__ with call-counting stub managers, not a local re-implementation of the decision. The first cut of this file modelled the decision inside the test; review correctly called that out as a model that can drift from the code it describes, and the second commit converts it. Two expectations written from reasoning turned out wrong and the real scheduler corrected them (every=1 starts with prefill, because the streak begins at 0).
• Cases covered: truth table over None/1/3/8, strict alternation at every=1, the latch case (decode due but unavailable -> prefills continue -> decode later fires exactly once), pure-decode fallback, no-work None path, a policy that exists but is disabled, decode-manager call counts (a disabled policy must not probe decode at all), and the config -> Scheduler.__init__ -> CLI wiring.
• Regression, tests/scheduler tests/kvcache tests/engine: 500 passed, 3 skipped. On unmodified main at 68a81ff the same command gives 470 passed; the delta is exactly the 30 new tests. The 2 failures (tests/engine/test_cache_budget.py, RuntimeError: Attention backend 'fi' requires flashinfer) are pre-existing and environmental — they fail identically on unmodified main, and this change touches neither that file nor flashinfer.

Notes for review

• One thing the deployment surfaced that is worth a reviewer's attention: the engine's Decode batch log line is throttled to every decode_log_interval (20 by default) decode steps, while Prefill batch logs every step. On a first pass this makes interleaving look like it is not working — 7 interleaved decode steps produce no log line at all. The verification above therefore measures the streamed inter-token gap, not the log.
• Interleaving and a larger prefill chunk are complementary, not substitutes: a larger chunk shortens the burst but still leaves decodes unscheduled for its whole duration.

Environment tested on

8x NVIDIA RTX 4090 24564 MiB, driver 595.58.03; 2x Intel Xeon Platinum 8468V; 1007 GiB RAM; Ubuntu 22.04.5, Python 3.10. Checkpoint: DeepSeek-V4-Flash (deepseek_v4, 43 layers, 256 routed experts, 1M native context). Serving command and the full measurement detail are in #483.

🤖 Generated with Claude Code

[huhoo]

Correction: the constants in this PR's description came from an older configuration

The PR body cites 384-token chunks, 0.388 s per chunk, and a 267 s stall. Those were measured when this deployment ran --swa-full-tokens-ratio 0.02. It now runs --swa-full-tokens-ratio 0.12 (308 window pages -> 3456-token chunk budget), so I re-measured under the configuration actually in use. Full detail in #483.

Re-measured, controlled, on the current config

Two requests, temperature: 0, no other traffic:

run	victim wall	victim output	decode rate
victim alone	11.8 s	699 tokens	59.4 tok/s
victim + one 200K-token prefill	90.4 s	699 tokens	7.7 tok/s

7.7x collapse. The log shows why:

14:03:48  DECODE  x19     6 s
14:03:56  prefill x58    77 s    <- 58 consecutive prefill steps; decode never scheduled
14:05:13  DECODE  x16     5 s

90.4 - 11.8 = 78.6 s added latency against a 77 s burst — the stall is total for its duration. The attacker prefilled 200,019 tokens in 58 chunks / 77 s (1.35 s per 3456-token chunk).

Corrected bound for this flag

The measured per-chunk cost is 1.35 s, not 0.388 s, so my earlier "~3 s at N=8" was optimistic by ~3.5x. Corrected (58 chunks / N decode steps, ~0.17 s per decode step at the measured single-stream rate):

--decode-interleave-every	worst-case stall	prefill added
unset	77 s (measured)	—
16	~21.6 s	~0.8%
8	~10.8 s	~1.6%
4	~5.4 s	~3.2%

The change itself is unaffected — the policy, the tests, and the zero-regression result all stand. What changes is the expected benefit: N=8 buys 77 s -> ~11 s, not 77 s -> 3 s. I would rather correct that here than let a reviewer size the win from a stale number.

One related win worth noting: the larger chunk also improved per-token prefill cost from 990 tok/s (384-token chunks) to 2598 tok/s (3456-token chunks), 2.6x. That is a separate effect and does not remove the starvation — it shortens the burst in seconds while leaving decode unscheduled for all of it.

Still not validated end-to-end on hardware: the flag has never been enabled on this deployment. The measurement above establishes the defect under the live config; the fix remains unit-tested only.