From 02653fa22cd1232cd252165b901091bd4a96f04f Mon Sep 17 00:00:00 2001 From: Thuongvu Ho Date: Wed, 14 Oct 2026 07:05:21 -0700 Subject: [PATCH] [Bugfix][KV Offload] Fix inaccurate proportional store dropping (#59140) Signed-off-by: Thuongvu Ho Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Or Ozeri --- .../kv_offload/tiering/test_backpressure.py | 111 +++++++++++++++++- vllm/v1/kv_offload/tiering/backpressure.py | 19 ++- 2 files changed, 120 insertions(+), 10 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_backpressure.py b/tests/v1/kv_offload/tiering/test_backpressure.py index b809bc4dc4..93717e8061 100644 --- a/tests/v1/kv_offload/tiering/test_backpressure.py +++ b/tests/v1/kv_offload/tiering/test_backpressure.py @@ -264,6 +264,41 @@ class TestBackpressure: if not tj.is_promotion and tj.submit_time > 0: tj.submit_time = now - age_s + def test_throttled_stores_update_metrics_and_release_pins(self, setup): + """Eight stores at a 75% drop rate yield two submissions and six drops.""" + bp = self.tier.bp_detector + assert isinstance(bp, EMABackpressureDetector) + bp._decay_half_life_s = 0.0 + bp._cooldown_s = 0.0 + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 2.5 + + all_keys = [] + next_id = 0 + for num_keys in [1, 2, 3, 1, 2, 3, 1, 2]: + keys = to_keys(range(next_id, next_id + num_keys)) + next_id += num_keys + all_keys.extend(keys) + self._store_blocks(keys) + + assert len(self.tier._held_jobs) == 2 + assert all( + self.primary._policy.get(key).ref_cnt == int(key in self.tier.blocks) + for key in all_keys + ) + + reduced = self.manager.get_stats().reduce() + stores_key = TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED + blocks_key = TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED + assert reduced["{}:('1:delayed',)".format(stores_key)] == 6 + assert reduced["{}:('1:delayed',)".format(blocks_key)] == 12 + + self._backdate_held_jobs(0.001) + self.tier.release_jobs() + self._simulate_on_schedule_end() + assert not self.manager._jobs + assert all(self.primary._policy.get(key).ref_cnt == 0 for key in all_keys) + def test_ema_updates_on_store_completion(self, setup): bp = self.tier.bp_detector fast_latency = 0.01 @@ -514,14 +549,23 @@ class TestThrottledDropPolicy: """Tests for proportional throttling under pressure.""" def _make_detector(self, **kwargs): + """Build a detector with cooldown and idle decay disabled.""" defaults = dict( high_water_s=_BP_HIGH_WATER_S, low_water_s=_BP_LOW_WATER_S, cooldown_s=0.0, + decay_half_life_s=0.0, ) defaults.update(kwargs) return EMABackpressureDetector(**defaults) + def _make_detector_at_rate(self, drop_rate: float): + """Build a pressured detector with the given drop rate.""" + bp = self._make_detector() + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * (1.0 + 2.0 * drop_rate) + return bp + def test_no_drop_when_not_under_pressure(self): bp = self._make_detector() for _ in range(20): @@ -574,6 +618,55 @@ class TestThrottledDropPolicy: for _ in range(20): assert bp.should_store(1) is False + @pytest.mark.parametrize("drop_rate", [0.0, 0.4, 0.75, 1.0]) + def test_drop_count_tracks_rate(self, drop_rate): + """After each decision, the drop count is within one store of the target.""" + bp = self._make_detector_at_rate(drop_rate) + drops = 0 + for num_stores in range(1, 21): + if not bp.should_store(1): + drops += 1 + expected_drops = num_stores * drop_rate + assert abs(drops - expected_drops) <= 1.0 + 1e-8 + + def test_drop_count_tracks_changing_rates(self): + """Dropped stores track the cumulative expected count as rates change.""" + bp = self._make_detector() + rates = [0.25] * 7 + [0.75] * 13 + [0.1] * 11 + drops = 0 + expected_drops = 0.0 + bp._under_pressure = True + for rate in rates: + bp.store_latency_ema = _BP_HIGH_WATER_S * (1.0 + 2.0 * rate) + drops += not bp.should_store(1) + expected_drops += rate + assert abs(drops - expected_drops) <= 1.0 + 1e-8 + + def test_fractional_state_survives_healthy_endpoints_and_metrics(self): + """Preserve the accumulated fraction until it contributes to a later drop.""" + bp = self._make_detector_at_rate(0.75) + assert bp.should_store(1) is True + bp.policy.pop_stores_dropped() + + bp._under_pressure = False + bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD + assert bp.is_healthy() + assert bp.should_store(1) is True + + bp._healthy_streak = 0 + assert not bp.is_healthy() + assert bp.should_store(1) is True + + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 3.0 + assert bp.should_store(1) is False + + bp.store_latency_ema = _BP_HIGH_WATER_S + assert bp.should_store(1) is True + + bp.store_latency_ema = _BP_HIGH_WATER_S * 2.5 + assert bp.should_store(1) is False + class TestSharedDropAccounting: """Both policies share drop accounting via DropAccountingPolicy.""" @@ -597,11 +690,19 @@ class TestSharedDropAccounting: policy.reset() assert policy.pop_stores_dropped() == (0, 0) - def test_throttled_reset_also_clears_call_count(self): - policy = ThrottledDropPolicy() - policy._call_count = 7 - policy.reset() - assert policy._call_count == 0 + def test_throttled_reset_restarts_decisions(self): + """Reset restarts the same deterministic sequence under steady pressure.""" + bp = EMABackpressureDetector( + high_water_s=_BP_HIGH_WATER_S, + low_water_s=_BP_LOW_WATER_S, + cooldown_s=0.0, + decay_half_life_s=0.0, + ) + bp._under_pressure = True + bp.store_latency_ema = _BP_HIGH_WATER_S * 2.5 + first = [bp.should_store(1) for _ in range(7)] + bp.policy.reset() + assert [bp.should_store(1) for _ in range(7)] == first class TestHealthyBypass: diff --git a/vllm/v1/kv_offload/tiering/backpressure.py b/vllm/v1/kv_offload/tiering/backpressure.py index deba37c94b..8214bb68ca 100644 --- a/vllm/v1/kv_offload/tiering/backpressure.py +++ b/vllm/v1/kv_offload/tiering/backpressure.py @@ -73,14 +73,20 @@ class ThrottledDropPolicy(DropAccountingPolicy): keeps secondary tiers populated under moderate pressure, reducing cache misses (and therefore recomputation / ITL regression) while still protecting TTFT under severe congestion. + + Accumulate ``drop_rate`` across store decisions. When the accumulator + reaches one, drop the store and subtract one from the accumulator, + carrying the remaining fraction into the next decision. """ def __init__(self, ramp_factor: float = 2.0): + """Initialize the drop accumulator and severity ramp.""" super().__init__() - self._call_count: int = 0 + self._drop_accumulator: float = 0.0 self._ramp_factor = ramp_factor def should_store(self, detector: BackpressureDetector) -> bool: + """Return whether to allow the store.""" # Proportional throttling needs the EMA and high watermark, which # only the EMA detector exposes. assert isinstance(detector, EMABackpressureDetector), ( @@ -98,13 +104,16 @@ class ThrottledDropPolicy(DropAccountingPolicy): return True if drop_rate >= 1.0: return False - self._call_count += 1 - period = max(2, round(1.0 / drop_rate)) - return (self._call_count % period) != 0 + self._drop_accumulator += drop_rate + if self._drop_accumulator >= 1.0: + self._drop_accumulator -= 1.0 + return False + return True def reset(self) -> None: + """Clear the drop accumulator and dropped store and block counts.""" super().reset() - self._call_count = 0 + self._drop_accumulator = 0.0 class BackpressureDetector(ABC):