diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 41b12bf40..ea0d4d0f5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -242,6 +242,23 @@ class Envs: # chunk still ends the scan by budget. See # docs_internal/perf/prefill-compute-intensity-plan.md S1. SGLANG_CP_PREFILL_MIX_CHUNKED = EnvBool(False) + # Cache-affinity batch formation (plan doc S4): keep FCFS, but skip a + # COLD candidate while the batch is warm-led so cache-hit requests group + # into one dense forward instead of being split by a cold 60K chunk. + # Bounded by the window/defer/age knobs below; disabled automatically + # under priority scheduling. + SGLANG_CP_PREFILL_AFFINITY_GROUP = EnvBool(False) + # Max cold candidates skipped per batch-formation pass (scan depth bound; + # each skip costs one re-match next pass). + SGLANG_CP_PREFILL_AFFINITY_WINDOW = EnvInt(16) + # Max consecutive passes the FCFS head may be deferred before the scan + # stops grouping and lets it lead the next (empty) batch. + SGLANG_CP_PREFILL_AFFINITY_MAX_DEFER = EnvInt(3) + # Absolute age bound (seconds) with the same effect as MAX_DEFER; must be + # well under SGLANG_REQ_WAITING_TIMEOUT when that is set. + SGLANG_CP_PREFILL_AFFINITY_MAX_AGE_S = EnvFloat(5.0) + # Warm/cold threshold in tokens on (device prefix hit + host hit). + SGLANG_CP_PREFILL_AFFINITY_WARM_FLOOR = EnvInt(64) # NSA MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM. # Lower values split query rows more aggressively to cap peak temporary memory. SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 6b6cbf0f6..b92c1502a 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -643,6 +643,9 @@ class Req(ReqDllmMixin): self.last_host_node: Any = None self.last_host_backup_node: Any = None self.host_hit_length = 0 + # Cache-affinity scheduling (plan doc S4): consecutive passes this + # request was deferred as the FCFS head while a warm batch formed. + self.affinity_defer_count = 0 self.prefix_match_deferred_by_pending_backup = False self.cp_hicache_prepared_backup = None # Tokens loaded from storage backend (L3) during prefetch for this request diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 6efdbfb46..c77bafb6c 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -101,6 +101,60 @@ def _cp_shared_kv_bs_gt1_admission_debug( logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args) +class AffinityDecision(Enum): + """Outcome of the cache-affinity scan decision (plan doc S4).""" + + ADMIT = auto() + SKIP_COLD = auto() + STOP = auto() + + +def decide_cp_prefill_affinity( + *, + is_warm: bool, + batch_empty_for_affinity: bool, + batch_warm_led: bool, + is_head: bool, + head_defer_count: int, + head_age_s: float, + window_used: int, +) -> AffinityDecision: + """Cache-affinity batch-formation decision (pure; plan doc S4, amended). + + The policy prevents exactly one thing: a COLD candidate joining a + WARM-led batch (one cold 60K extend turns a 1-2s warm forward into a + 5-10s one). Everything else admits: + + - WARM always admits — into a warm-led batch it grows the dense group; + into a cold-led batch it is free density (the cold extend dominates + the forward time anyway). + - COLD into an empty batch admits — the FCFS head always starts a + batch, so the queue keeps moving and a deferred cold leads the very + next batch. + - COLD into a COLD-led batch admits — small cold requests co-batch + today (the CP caps bound them); refusing would be a regression. + - COLD into a WARM-led batch is skipped, bounded three ways: the + per-pass window, and for the FCFS head a defer count and an age + bound. On any bound tripping the decision is STOP (end the warm + batch cleanly) rather than admitting the cold into it: the cold + waits for the same forward either way, but leads its own clean batch + next pass instead of polluting this one. + """ + + if is_warm: + return AffinityDecision.ADMIT + if batch_empty_for_affinity or not batch_warm_led: + return AffinityDecision.ADMIT + if window_used >= int(envs.SGLANG_CP_PREFILL_AFFINITY_WINDOW.get()): + return AffinityDecision.STOP + if is_head and ( + head_defer_count >= int(envs.SGLANG_CP_PREFILL_AFFINITY_MAX_DEFER.get()) + or head_age_s >= float(envs.SGLANG_CP_PREFILL_AFFINITY_MAX_AGE_S.get()) + ): + return AffinityDecision.STOP + return AffinityDecision.SKIP_COLD + + class CacheAwarePolicy(Enum): """Scheduling policies that are aware of the tree cache.""" diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index eba94e203..d970e08bb 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -69,6 +69,7 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r from sglang.srt.layers.attention.mamba.ops import ( initialize_mamba_selective_state_update_backend, ) +from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split from sglang.srt.layers.dp_attention import ( compute_dp_attention_world_info, get_attention_cp_group, @@ -163,8 +164,10 @@ from sglang.srt.managers.scheduler_graceful_shutdown_mixin import ( ) from sglang.srt.managers.schedule_policy import ( AddReqResult, + AffinityDecision, PrefillAdder, SchedulePolicy, + decide_cp_prefill_affinity, ) from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker @@ -2562,6 +2565,20 @@ class Scheduler( running_loras = {req.lora_id for req in self.running_batch.reqs} # Get requests from the waiting queue to a new prefill batch + # Cache-affinity grouping (plan doc S4): active only for the CP + # shared-KV bs>1 prefill context, and never under priority + # scheduling (the skip must not reorder across priority classes — + # disabled wholesale as the conservative guard). + affinity_on = ( + envs.SGLANG_CP_PREFILL_AFFINITY_GROUP.get() + and not self.enable_priority_scheduling + and self.server_args.enable_cp_shared_kv_prefill_bs_gt1 + and is_nsa_prefill_cp_in_seq_split() + ) + affinity_window_used = 0 + affinity_head_pending = True # first classified candidate = FCFS head + affinity_batch_warm_led: Optional[bool] = None + for req in self.waiting_queue: if self.enable_lora and req.lora_id not in running_loras: if self.enable_lora_overlap_loading: @@ -2606,11 +2623,54 @@ class Scheduler( ) req.init_next_round_input(self.tree_cache) + + if affinity_on: + # Decision strictly post-match / pre-admit: a skipped + # candidate has no lock, no allocation, no budget mutation + # to unwind, and re-matching it next pass is exactly what + # the scan does today after a cap rejection. + req_is_warm = ( + len(req.prefix_indices) + + int(getattr(req, "host_hit_length", 0) or 0) + ) >= int(envs.SGLANG_CP_PREFILL_AFFINITY_WARM_FLOOR.get()) + is_head = affinity_head_pending + affinity_head_pending = False + decision = decide_cp_prefill_affinity( + is_warm=req_is_warm, + batch_empty_for_affinity=( + len(adder.can_run_list) + - (1 if adder._chunked_req_in_batch is not None else 0) + == 0 + ), + batch_warm_led=bool(affinity_batch_warm_led), + is_head=is_head, + head_defer_count=req.affinity_defer_count, + head_age_s=( + time.perf_counter() + - req.time_stats.wait_queue_entry_time + if req.time_stats.wait_queue_entry_time > 0.0 + else 0.0 + ), + window_used=affinity_window_used, + ) + if decision is AffinityDecision.SKIP_COLD: + if is_head: + req.affinity_defer_count += 1 + affinity_window_used += 1 + continue + if decision is AffinityDecision.STOP: + break + + num_admitted_before = len(adder.can_run_list) res = adder.add_one_req( req, has_chunked_req=(self.chunked_req is not None), truncation_align_size=self.truncation_align_size, ) + if affinity_on and len(adder.can_run_list) > num_admitted_before: + req.affinity_defer_count = 0 + if affinity_batch_warm_led is None: + affinity_batch_warm_led = req_is_warm if self.enable_lora: running_loras.add(req.lora_id) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index e28547c05..2d6914f4f 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -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.