Group cache-hit prefills into dense batches (SGLANG_CP_PREFILL_AFFINITY_GROUP)

专题 S4 (design: docs_internal/perf/prefill-compute-intensity-plan.md S4,
amended).  Under FCFS a cold request joining a warm-led batch turns a
1-2s cache-hit forward into a 5-10s one, splitting the warm work into
the 新-cache-新 pattern.  The policy prevents exactly that one thing:

- WARM candidates always admit (into a cold-led batch they are free
  density — the cold extend dominates the forward anyway).
- COLD admits into an empty or cold-led batch (small colds co-batch
  today; the FCFS head always starts a batch so the queue keeps moving).
- COLD into a WARM-led batch is skipped, bounded by a per-pass window
  (W=16 skips), a head defer count (K=3 passes) and an age bound
  (T=5s).  On any bound the scan STOPS instead of force-admitting: the
  cold waits for the same forward either way, but leads its own clean
  batch next pass instead of polluting this one.

The skip is strictly post-match / pre-admit (after init_next_round_input,
before add_one_req): no lock, no allocation, no budget mutation to
unwind, and re-matching a skipped candidate next pass is exactly what
the scan already does after a cap rejection.  Classification is the
in-scan match result (device prefix + host hit vs a 64-token floor) —
under FCFS+L2 no pre-scan signal exists, so this adds zero matching
work for inspected candidates.  Disabled wholesale under priority
scheduling (the skip must not reorder across priority classes).

Three amendments vs the design draft, reasoned in the decision-table
docstring: cold+cold-led admits (STOP would regress today's small-cold
co-batching); starved heads STOP rather than force-admit (clean batch
boundaries at identical latency); priority interaction handled by
disabling rather than per-request comparison.

Decision logic is a pure function with table + bounds unit tests
(28/28 adder suite green).  Default OFF.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-06-12 07:44:39 +00:00
co-authored by Claude Fable 5
parent 54c056af83
commit 6a1e862f48
5 changed files with 248 additions and 0 deletions
@@ -847,6 +847,120 @@ class TestPrefillAdder(CustomTestCase):
)
self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"])
def test_affinity_decision_table(self):
# Plan doc S4 (amended): the policy prevents exactly one thing — a
# COLD candidate joining a WARM-led batch. Everything else admits.
from sglang.srt.managers.schedule_policy import (
AffinityDecision,
decide_cp_prefill_affinity,
)
base = dict(
is_head=False,
head_defer_count=0,
head_age_s=0.0,
window_used=0,
)
# WARM always admits, whatever the batch looks like.
for empty, warm_led in [(True, False), (False, True), (False, False)]:
self.assertIs(
decide_cp_prefill_affinity(
is_warm=True,
batch_empty_for_affinity=empty,
batch_warm_led=warm_led,
**base,
),
AffinityDecision.ADMIT,
)
# COLD into an empty batch admits (the head always starts a batch).
self.assertIs(
decide_cp_prefill_affinity(
is_warm=False,
batch_empty_for_affinity=True,
batch_warm_led=False,
**base,
),
AffinityDecision.ADMIT,
)
# COLD into a COLD-led batch admits (small colds co-batch today).
self.assertIs(
decide_cp_prefill_affinity(
is_warm=False,
batch_empty_for_affinity=False,
batch_warm_led=False,
**base,
),
AffinityDecision.ADMIT,
)
# COLD into a WARM-led batch is skipped (within bounds).
self.assertIs(
decide_cp_prefill_affinity(
is_warm=False,
batch_empty_for_affinity=False,
batch_warm_led=True,
**base,
),
AffinityDecision.SKIP_COLD,
)
def test_affinity_anti_starvation_bounds_stop_the_scan(self):
from sglang.srt.managers.schedule_policy import (
AffinityDecision,
decide_cp_prefill_affinity,
)
cold_in_warm = dict(
is_warm=False,
batch_empty_for_affinity=False,
batch_warm_led=True,
)
# Head deferred K times -> STOP (end the warm batch cleanly; the
# cold leads the next, empty batch — never force-polluted into this
# one).
self.assertIs(
decide_cp_prefill_affinity(
**cold_in_warm,
is_head=True,
head_defer_count=3,
head_age_s=0.0,
window_used=0,
),
AffinityDecision.STOP,
)
# Head older than T -> STOP regardless of defer count.
self.assertIs(
decide_cp_prefill_affinity(
**cold_in_warm,
is_head=True,
head_defer_count=0,
head_age_s=10.0,
window_used=0,
),
AffinityDecision.STOP,
)
# Window exhausted -> STOP, head or not.
self.assertIs(
decide_cp_prefill_affinity(
**cold_in_warm,
is_head=False,
head_defer_count=0,
head_age_s=0.0,
window_used=16,
),
AffinityDecision.STOP,
)
# A NON-head cold within bounds is skipped without K/T applying.
self.assertIs(
decide_cp_prefill_affinity(
**cold_in_warm,
is_head=False,
head_defer_count=99,
head_age_s=99.0,
window_used=0,
),
AffinityDecision.SKIP_COLD,
)
def test_add_chunked_req_seeds_true_prefix_into_cp_budget(self):
# C1 (plan doc S1.1-1a): the chunk's carried prefix must count toward
# the CP cached tally so later admission gates see its footprint.